Upload app.py

app.py

@@ -88,11 +88,8 @@ def create_refseg_demo(model, tokenizer, device):
     
     def on_submit(image, text):
         # Convert PIL -> np array
-        image_np = np.array(image).copy()
         transform = transforms.ToTensor()
         image_t = transform(image).unsqueeze(0).to(device)
-        image_t = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])(image_t)
-        image_t = torch.nn.functional.interpolate(image_t, (512, 512), mode='bilinear', align_corners=True)
 
         with torch.no_grad():
             out = model(image_t, text)
@@ -103,24 +100,23 @@ def create_refseg_demo(model, tokenizer, device):
         else:
             mask = out
 
-        # Convert to binary mask
+        # If model returns multi-channel, collapse with argmax
         if mask.ndim > 2:
             mask = np.argmax(mask, axis=0)
-        mask = (mask > 0).astype(np.uint8)
 
-        # Resize mask to original image size
-        mask = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+        mask = mask.astype(np.uint8)
 
-        # Overlay mask
+        # Overlay mask on original image
+        image_np = np.array(image).copy()
         alpha = 0.65
-        overlay = image_np.copy()
-        overlay[mask == 0] = (overlay[mask == 0] * alpha).astype(np.uint8)
+        image_np[mask == 0] = (image_np[mask == 0] * alpha).astype(np.uint8)
 
         # Draw contours
         contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(overlay, contours, -1, (0, 255, 0), 2)
+        cv2.drawContours(image_np, contours, -1, (0, 255, 0), 2)
+
+        return Image.fromarray(image_np)
 
-        return Image.fromarray(overlay)
 
     submit.click(on_submit, inputs=[input_image, input_text], outputs=refseg_image)
     examples = gr.Examples(