Update app.py

app.py

@@ -6,40 +6,22 @@ from typing import Union
 
 class Preprocessor:
     def __init__(self):
-        """
-        Initialize the preprocessing transformations.
-        """
         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:
-        """
-        Apply preprocessing to the input image.
-
-        :param image: Input image to be preprocessed.
-        :return: Preprocessed image as a tensor.
-        """
         return self.transform(image)
 
 class SegmentationModel:
     def __init__(self):
-        """
-        Initialize and load the DeepLabV3 ResNet101 model.
-        """
         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:
-        """
-        Perform inference using the model on the input batch.
-
-        :param input_batch: Batch of preprocessed images.
-        :return: Model output tensor.
-        """
         with torch.no_grad():
             if torch.cuda.is_available():
                 input_batch = input_batch.to('cuda')
@@ -48,40 +30,22 @@ class SegmentationModel:
 
 class OutputColorizer:
     def __init__(self):
-        """
-        Initialize the color palette for segmentations.
-        """
         palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
-        colors : torch.Tensor = torch.as_tensor([i for i in range(21)])[:, None] * palette
+        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:
-        """
-        Apply colorization to the segmentation output.
-
-        :param output: Segmentation output tensor.
-        :return: Colorized segmentation 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):
-        """
-        Initialize the Segmenter with Preprocessor, SegmentationModel, and OutputColorizer.
-        """
         self.preprocessor = Preprocessor()
         self.model = SegmentationModel()
         self.colorizer = OutputColorizer()
 
     def segment(self, image: Union[Image.Image, torch.Tensor]) -> Image.Image:
-        """
-        Perform the complete segmentation process on the input image.
-
-        :param image: Input image to be segmented.
-        :return: Colorized segmentation image.
-        """
         input_image: Image.Image = image.convert("RGB")
         input_tensor: torch.Tensor = self.preprocessor(input_image)
         input_batch: torch.Tensor = input_tensor.unsqueeze(0)
@@ -99,4 +63,5 @@ interface = gr.Interface(
     description="Upload an image to perform semantic segmentation using Deeplabv3 ResNet101."
 )
 
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()