Upload 2 files

app.py

@@ -0,0 +1,78 @@
+import torch
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import gradio as gr
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+
+# Load processor and model from Hugging Face
+processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-coco-instance")
+model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-small-coco-instance")
+model.eval()
+
+# Load label map from model config
+COCO_INSTANCE_CATEGORY_NAMES = model.config.id2label if hasattr(model.config, "id2label") else [str(i) for i in range(133)]
+
+def segment_image(image, threshold=0.5):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    results = processor.post_process_instance_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+
+    segmentation_map = results["segmentation"].cpu().numpy()  # shape: [H, W]
+    segments_info = results["segments_info"]  # list of dicts with keys: id, label_id, score
+
+    image_np = np.array(image).copy()
+    overlay = image_np.copy()
+    fig, ax = plt.subplots(1, figsize=(10, 10))
+    ax.imshow(image_np)
+
+    for segment in segments_info:
+        score = segment.get("score", 1.0)
+        if score < threshold:
+            continue
+
+        segment_id = segment["id"]
+        label_id = segment["label_id"]
+        mask = segmentation_map == segment_id
+
+        # Random color per object
+        color = np.random.rand(3)
+        overlay[mask] = (overlay[mask] * 0.5 + np.array(color) * 255 * 0.5).astype(np.uint8)
+
+        # Draw bounding box
+        y_indices, x_indices = np.where(mask)
+        if len(x_indices) == 0 or len(y_indices) == 0:
+            continue
+        x1, x2 = x_indices.min(), x_indices.max()
+        y1, y2 = y_indices.min(), y_indices.max()
+
+        label_name = COCO_INSTANCE_CATEGORY_NAMES.get(str(label_id), str(label_id))
+        ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, color=color, linewidth=2))
+        ax.text(x1, y1, f"{label_name}: {score:.2f}",
+                bbox=dict(facecolor='yellow', alpha=0.5), fontsize=10)
+
+    ax.imshow(overlay)
+    ax.axis('off')
+    output_path = "mask2former_output.png"
+    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
+    plt.close()
+    return output_path
+
+
+
+# Gradio interface
+interface = gr.Interface(
+    fn=segment_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Confidence Threshold")
+    ],
+    outputs=gr.Image(type="filepath", label="Segmented Output"),
+    title="Mask2Former Instance Segmentation (Transformer)",
+    description="Upload an image to segment objects using Facebook's transformer-based Mask2Former model (Swin-Small backbone)."
+)
+
+if __name__ == "__main__":
+    interface.launch(debug=True,share=True)

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0.0
+transformers>=4.40.0
+gradio>=4.24.0
+matplotlib
+Pillow
+numpy
+scipy