initial commit

ex6.py

@@ -0,0 +1,57 @@
+import gradio as gr
+from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
+from PIL import Image
+import numpy as np
+
+feature_extractor = MobileViTFeatureExtractor.from_pretrained("apple/deeplabv3-mobilevit-small")
+model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-small")
+
+#(21 classes)
+COLORS = np.array([
+    [0, 0, 0],      
+    [128, 0, 0], 
+    [0, 128, 0],    
+    [128, 128, 0],  
+    [0, 0, 128],   
+    [128, 0, 128],  
+    [0, 128, 128],  
+    [128, 128, 128],
+    [64, 0, 0],     
+    [192, 0, 0],    
+    [64, 128, 0],   
+    [192, 128, 0], 
+    [64, 0, 128],   
+    [192, 0, 128],  
+    [64, 128, 128], 
+    [192, 128, 128],
+    [0, 64, 0],     
+    [128, 64, 0],   
+    [0, 192, 0],    
+    [128, 192, 0],  
+    [0, 64, 128],   
+    [128, 64, 128]  
+], dtype=np.uint8)  # Ensure the data type is uint8 for image processing
+
+def segment_image(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_mask = logits.argmax(1).squeeze(0).numpy()
+    
+    colored_mask = COLORS[predicted_mask]
+    colored_mask_image = Image.fromarray(colored_mask)
+    colored_mask_resized = colored_mask_image.resize(image.size, Image.NEAREST)
+    
+    return colored_mask_resized
+
+interface = gr.Interface(
+    fn=segment_image,  
+    inputs=gr.Image(type="pil"),  
+    outputs="image",  
+    title="Image Segmentation with MobileViT",
+    description="Upload an image to see the semantic segmentation result. The segmentation mask uses different colors to indicate different classes.",
+)
+
+interface.launch(share=True)
+