Update app.py

app.py

@@ -1,80 +1,82 @@
 import torch
 from torchvision import models, transforms
-from PIL import Image
+from PIL import Image, ImageDraw
 import gradio as gr
 from typing import Union
 
+SEGMENTATION_MODELS = {
+    "deeplabv3_resnet101": models.segmentation.deeplabv3_resnet101,}
+
+class ModelLoader:
+    def __init__(self, model_dict: dict):
+        self.model_dict = model_dict
+    
+    def load_model(self, model_name: str) -> torch.nn.Module:
+        model_name_lower = model_name.lower()
+        if model_name_lower in self.model_dict:
+            model_class = self.model_dict[model_name_lower]
+            model = model_class(pretrained=True)
+            model.eval()
+            return model
+        else:
+            raise ValueError(f"Model {model_name} is not supported")
+
 class Preprocessor:
-    def __init__(self):
-        self.transform = transforms.Compose([
+    def __init__(self, transform: transforms.Compose = 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
+        ])):
+        self.transform = transform
+    
+    def preprocess(self, image: Image.Image) -> torch.Tensor:
+        return self.transform(image).unsqueeze(0)
 
-class OutputColorizer:
+class Postprocessor:
     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')
+    
+    def postprocess(self, output: torch.Tensor) -> Image.Image:
+        output_predictions = output.argmax(0)
+        colorized_output = Image.fromarray(output_predictions.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()
+class Segmentation:
+    def __init__(self, model_loader: ModelLoader, preprocessor: Preprocessor, postprocessor: Postprocessor):
+        self.model_loader = model_loader
+        self.preprocessor = preprocessor
+        self.postprocessor = postprocessor
+    
+    def segment(self, image: Image.Image, selected_model: str) -> Image.Image:
+        model = self.model_loader.load_model(selected_model)
+        input_tensor = self.preprocessor.preprocess(image)
+
+        if torch.cuda.is_available():
+            input_tensor = input_tensor.to("cuda")
+            model = model.to("cuda")      
 
-    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)
+        with torch.no_grad():
+            output = model(input_tensor)['out'][0]
+        return self.postprocessor.postprocess(output)
 
 class GradioApp:
-    def __init__(self, segmenter: Segmenter):
-        self.segmenter = segmenter
-
+    def __init__(self, segmentation: Segmentation):
+        self.segmentation = segmentation
+    
     def launch(self):
         with gr.Blocks() as demo:
             gr.Markdown("<h1 style='text-align: center; color: #4CAF50;'>Deeplabv3 Segmentation</h1>")
             gr.Markdown("<p style='text-align: center;'>Upload an image to perform semantic segmentation using Deeplabv3 ResNet101.</p>")
-            gr.Markdown("""
-                ### Model Information
-                **DeepLabv3 with ResNet101** is a convolutional neural network model designed for semantic image segmentation. 
-                It utilizes atrous convolution to capture multi-scale context by using different atrous rates.
-                """)
             with gr.Row():
                 with gr.Column():
-                    image_input = gr.Image(type='pil', label="Input Image", show_label=False)
+                    upload_image = gr.Image(type='pil', label="Upload Image")
+                    self.model_dropdown = gr.Dropdown(choices=list(SEGMENTATION_MODELS.keys()), label="Select Model")
+                    segment_button = gr.Button("Segment")
                 with gr.Column():
-                    image_output = gr.Image(type='pil', label="Segmented Output", show_label=False)
-                    
-            button = gr.Button("Segment")
-            button.click(fn=self.segmenter.segment, inputs=image_input, outputs=image_output)
-
+                    output_image = gr.Image(type='pil', label="Segmented Output")
+            segment_button.click(fn=self.segmentation.segment, inputs=[upload_image, self.model_dropdown], outputs=output_image)
             gr.Markdown("### Example Images")
             gr.Examples(
                 examples=[
@@ -82,14 +84,16 @@ class GradioApp:
                     ["https://www.timeforkids.com/wp-content/uploads/2023/09/G3G5_230915_puffins_on_the_rise.jpg?w=1024"],
                     ["https://www.timeforkids.com/wp-content/uploads/2024/03/G3G5_240412_bug_eyed.jpg?w=1024"]
                 ],
-                inputs=image_input,
-                outputs=image_output,
+                inputs=upload_image,
+                outputs=output_image,
                 label="Click an example to use it"
             )
-            
         demo.launch()
 
 if __name__ == "__main__":
-    segmenter = Segmenter()
-    app = GradioApp(segmenter)
+    model_loader = ModelLoader(SEGMENTATION_MODELS)
+    preprocessor = Preprocessor()
+    postprocessor = Postprocessor()
+    segmentation = Segmentation(model_loader, preprocessor, postprocessor)
+    app = GradioApp(segmentation)
     app.launch()