Update app.py

app.py

@@ -1,15 +1,16 @@
 import numpy as np
 from PIL import Image
 import torch
-from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
+from transformers import AutoImageProcessor, MobileViTForSemanticSegmentation
 import gradio as gr
 
 
 # ---------------------------
-# Load model & preprocessing
+# Load model & processor
 # ---------------------------
 model_checkpoint = "apple/deeplabv3-mobilevit-small"
-feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_checkpoint)
+
+image_processor = AutoImageProcessor.from_pretrained(model_checkpoint)
 model = MobileViTForSemanticSegmentation.from_pretrained(model_checkpoint).eval()
 
 palette = np.array(
@@ -39,27 +40,30 @@ def predict(image):
         return None, None
 
     with torch.no_grad():
-        inputs = feature_extractor(image, return_tensors="pt")
+        inputs = image_processor(image, return_tensors="pt")
         outputs = model(**inputs)
 
-    # Convert back to uint8 image (resized/cropped)
+    # Re-normalize back to uint8
     resized = (
-        inputs["pixel_values"].numpy().squeeze().transpose(1, 2, 0)[..., ::-1] * 255
+        inputs["pixel_values"]
+        .numpy()
+        .squeeze()
+        .transpose(1, 2, 0)[..., ::-1] * 255
     ).astype(np.uint8)
 
-    # Segmentation classes
+    # Class map
     classes = outputs.logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
 
-    # Vectorized coloring (FAST)
+    # Vectorized lookup table coloring
     colored = palette[classes]
 
-    # Resize seg mask to match resized input resolution
+    # Resize segmentation to match resized input
     colored_img = Image.fromarray(colored).resize(
         (resized.shape[1], resized.shape[0]),
         resample=Image.Resampling.NEAREST
     )
 
-    # Mask: everything except background
+    # Binary mask for overlay
     mask = (classes != 0).astype(np.uint8) * 255
     mask_img = Image.fromarray(mask).resize(
         (resized.shape[1], resized.shape[0]),
@@ -73,9 +77,10 @@ def predict(image):
 
 
 # ---------------------------
-# Label HTML
+# Labels HTML
 # ---------------------------
-inverted = {0,1,4,5,8,9,12,13,16,17,20}
+inverted = {0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20}
+
 labels_html = " ".join(
     f"<span style='background-color: rgb{tuple(palette[i])}; "
     f"color: {'white' if i in inverted else 'black'}; padding: 2px 4px;'>"
@@ -100,7 +105,7 @@ article = """
 
 
 # ---------------------------
-# Modern Gradio App (Blocks)
+# Gradio App (Blocks)
 # ---------------------------
 with gr.Blocks(title="Semantic Segmentation with MobileViT") as demo:
     gr.Markdown("# Semantic Segmentation with MobileViT & DeepLabV3")