Update app.py

app.py

@@ -15,7 +15,7 @@ sam_processor = SamProcessor.from_pretrained("jadechoghari/robustsam-vit-base")
 
 def apply_mask(image, mask, color):
     """Apply a mask to an image with a specific color."""
-    for c in range(3):  # iterate over rgb channels
+    for c in range(3):  # Iterate over RGB channels
         image[:, :, c] = np.where(mask, color[c], image[:, :, c])
     return image
 
@@ -59,33 +59,42 @@ def query(image, texts, threshold):
                 inputs["reshaped_input_sizes"].cpu()
             )[0][0][0].numpy()
             
-            # we apply the mask with the corresponding color
-            color = colors[i % len(colors)]  # we cycle through colors
+            color = colors[i % len(colors)]  # cycle through colors
             image = apply_mask(image, mask > 0.5, color)
 
     result_image = Image.fromarray(image)
     
     return result_image
 
-description = (
-    "Welcome to RobustSAM by Snap Research."
-    "This Space uses RobustSAM, a robust version of the Segment Anything Model (SAM) with improved performance on low-quality images while maintaining zero-shot segmentation capabilities. "
-    "Thanks to its integration with OWLv2, RobustSAM becomes text-promptable, allowing for flexible and accurate segmentation, even with degraded image quality. Try the example or input an image with comma-separated candidate labels to see the enhanced segmentation results."
-)
-
-demo = gr.Interface(
-    query,
-    inputs=[gr.Image(type="pil", label="Image Input"), gr.Textbox(label="Candidate Labels"), gr.Slider(0, 1, value=0.05, label="Confidence Threshold")],
-    outputs=gr.Image(type="pil", label="Segmented Image"),
-    title="RobustSAM",
-    description=description,
-    examples=[
+title = """
+# RobustSAM
+"""
+
+description = """
+**Welcome to RobustSAM by Snap Research.**
+
+This Space uses **RobustSAM**, a robust version of the Segment Anything Model (SAM) with improved performance on low-quality images while maintaining zero-shot segmentation capabilities.
+
+Thanks to its integration with **OWLv2**, RobustSAM becomes text-promptable, allowing for flexible and accurate segmentation, even with degraded image quality.
+
+Try the example or input an image with comma-separated candidate labels to see the enhanced segmentation results.
+
+For better results, please check the [GitHub repository](https://github.com/robustsam/RobustSAM).
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    
+    gr.Interface(
+        query,
+        inputs=[gr.Image(type="pil", label="Image Input"), gr.Textbox(label="Candidate Labels"), gr.Slider(0, 1, value=0.05, label="Confidence Threshold")],
+        outputs=gr.Image(type="pil", label="Segmented Image"),
+        examples=[
         ["./blur.jpg", "insect", 0.1],
-        ["./lowlight.jpg", "bus, window", 0.1],
-        ["./rain.jpg", "tree, leafs", 0.1],
-        ["./haze.jpg", "", 0.1],
-    ],
-    cache_examples=True
-)
-
-demo.launch()
+        ["./lowlight.jpg", "bus, window", 0.1]
+        ],
+        cache_examples=True
+    )
+
+demo.launch()