| import torch | |
| from torchvision import models, transforms | |
| from PIL import Image | |
| import gradio as gr | |
| from typing import Union | |
| class Preprocessor: | |
| def __init__(self): | |
| self.transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def __call__(self, image: Image.Image) -> torch.Tensor: | |
| return self.transform(image) | |
| class SegmentationModel: | |
| def __init__(self): | |
| self.model = models.segmentation.deeplabv3_resnet101(pretrained=True) | |
| self.model.eval() | |
| if torch.cuda.is_available(): | |
| self.model.to('cuda') | |
| def predict(self, input_batch: torch.Tensor) -> torch.Tensor: | |
| with torch.no_grad(): | |
| if torch.cuda.is_available(): | |
| input_batch = input_batch.to('cuda') | |
| output: torch.Tensor = self.model(input_batch)['out'][0] | |
| return output | |
| class OutputColorizer: | |
| def __init__(self): | |
| palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1]) | |
| colors = torch.as_tensor([i for i in range(21)])[:, None] * palette | |
| self.colors = (colors % 255).numpy().astype("uint8") | |
| def colorize(self, output: torch.Tensor) -> Image.Image: | |
| colorized_output = Image.fromarray(output.byte().cpu().numpy(), mode='P') | |
| colorized_output.putpalette(self.colors.ravel()) | |
| return colorized_output | |
| class Segmenter: | |
| def __init__(self): | |
| self.preprocessor = Preprocessor() | |
| self.model = SegmentationModel() | |
| self.colorizer = OutputColorizer() | |
| def segment(self, image: Union[Image.Image, torch.Tensor]) -> Image.Image: | |
| input_image: Image.Image = image.convert("RGB") | |
| input_tensor: torch.Tensor = self.preprocessor(input_image) | |
| input_batch: torch.Tensor = input_tensor.unsqueeze(0) | |
| output: torch.Tensor = self.model.predict(input_batch) | |
| output_predictions: torch.Tensor = output.argmax(0) | |
| return self.colorizer.colorize(output_predictions) | |
| segmenter = Segmenter() | |
| interface = gr.Interface( | |
| fn=segmenter.segment, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Image(type="pil"), | |
| title="Deeplabv3 Segmentation", | |
| description="Upload an image to perform semantic segmentation using Deeplabv3 ResNet101." | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |