Upload 2 files

app.py

@@ -0,0 +1,101 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from torchvision import models, transforms
+from PIL import Image
+from rankseg import RankSEG
+import numpy as np
+
+# 1. Load Model (Cache it to avoid reloading)
+# We use a lightweight DeepLabV3+ (MobileNet) for speed in the demo, 
+# or ResNet50 if we want better quality. Let's use ResNet50.
+def load_model():
+    try:
+        weights = models.segmentation.DeepLabV3_ResNet50_Weights.DEFAULT
+        model = models.segmentation.deeplabv3_resnet50(weights=weights)
+    except:
+        model = models.segmentation.deeplabv3_resnet50(pretrained=True)
+    model.eval()
+    return model
+
+model = load_model()
+
+# 2. Define Transformations
+preprocess = transforms.Compose([
+    transforms.Resize(520),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Color palette for visualization (PASCAL VOC style)
+def get_palette():
+    return [
+        (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)
+    ]
+
+def colorize_mask(mask):
+    # mask: (H, W) numpy array
+    palette = get_palette()
+    h, w = mask.shape
+    color_mask = np.zeros((h, w, 3), dtype=np.uint8)
+    for label, color in enumerate(palette):
+        color_mask[mask == label] = color
+    return color_mask
+
+# 3. Inference Function
+def predict(image):
+    if image is None:
+        return None, None
+    
+    # Preprocess
+    input_tensor = preprocess(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = model(input_tensor)['out']
+        probs = F.softmax(output, dim=1) # (1, 21, H, W)
+        
+    # --- METHOD 1: ARGMAX ---
+    argmax_pred = torch.argmax(probs, dim=1).squeeze().numpy()
+    argmax_vis = colorize_mask(argmax_pred)
+    
+    # --- METHOD 2: RANKSEG ---
+    # Optimize for Dice
+    rankseg = RankSEG(metric='dice', solver='RMA')
+    rankseg_pred_tensor = rankseg.predict(probs)
+    rankseg_pred = rankseg_pred_tensor.squeeze().numpy()
+    rankseg_vis = colorize_mask(rankseg_pred)
+    
+    return argmax_vis, rankseg_vis
+
+# 4. Gradio Interface
+title = "🧩 RankSEG: Optimize Segmentation Metrics without Retraining"
+description = """
+**RankSEG** (NeurIPS 2025) is a plug-and-play module that improves segmentation results by directly optimizing for Dice/IoU metrics during inference.
+Upload an image to see how RankSEG refines the mask compared to standard Argmax.
+"""
+
+# examples = [
+#     ["example1.jpg"], # You will need to upload example images to the HF Space
+#     ["example2.jpg"]
+# ]
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=[
+        gr.Image(label="Standard Argmax Prediction"),
+        gr.Image(label="RankSEG Optimized (Dice)")
+    ],
+    title=title,
+    description=description,
+    # examples=examples, # Uncomment if you have examples
+    cache_examples=False
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.0.0
+torch
+torchvision
+pillow
+numpy
+rankseg