Update

app.py

@@ -9,8 +9,9 @@ import os
 import pandas as pd
 from utils import *
 from PIL import Image
+from gradio_image_prompter import ImagePrompter
+
 
-# Carga de modelos
 sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-huge")
 sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-huge")
 
@@ -31,23 +32,26 @@ def predict_masks_and_scores(model, processor, raw_image, input_points=None, inp
     scores = outputs.iou_scores
     return masks, scores
 
-def encode_pil_to_base64(pil_image):
-    buffer = BytesIO()
-    pil_image.save(buffer, format="PNG")
-    return base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-def compare_images_points_and_masks(user_image, input_boxes, input_points):
-    for example_path, example_data in example_data_map.items():
-        if example_data["size"] == list(user_image.size):
-            user_image = Image.open(example_data['original_image_path'])
-    input_boxes = input_boxes.values.tolist()
-    input_points = input_points.values.tolist()
-
-    input_boxes = [[[int(coord) for coord in box] for box in input_boxes if any(box)]]
-    input_points = [[[int(coord) for coord in point] for point in input_points if any(point)]]
-
-    input_boxes = input_boxes if input_boxes[0] else None
-    input_points = input_points if input_points[0] else None
+def process_inputs(prompts):
+    raw_entries = prompts["points"]
+    
+    input_points = []
+    input_boxes = []
+
+    for entry in raw_entries:
+        x1, y1, type_, x2, y2, cls = entry
+        if type_ == 1:
+            input_points.append([int(x1), int(y1)])
+        elif type_ == 2:
+            x_min = int(min(x1, x2))
+            y_min = int(min(y1, y2))
+            x_max = int(max(x1, x2))
+            y_max = int(max(y1, y2))
+            input_boxes.append([x_min, y_min, x_max, y_max])
+
+    input_boxes = [input_boxes] if input_boxes else None
+    input_points = [input_points] if input_points else None
+    user_image = prompts['image']
 
     sam_masks, sam_scores = predict_masks_and_scores(sam_model, sam_processor, user_image, input_boxes=input_boxes, input_points=input_points)
     sam_hq_masks, sam_hq_scores = predict_masks_and_scores(sam_hq_model, sam_hq_processor, user_image, input_boxes=input_boxes, input_points=input_points)
@@ -89,22 +93,7 @@ def compare_images_points_and_masks(user_image, input_boxes, input_points):
     """
     return html_code
 
-def load_examples(json_file="examples.json"):
-    with open(json_file, "r") as f:
-        examples = json.load(f)
-    return examples
-
-examples = load_examples()
-example_paths = [example["image_path"] for example in examples]
-example_data_map = {
-    example["image_path"]: {
-        "original_image_path": example["original_image_path"],
-        "points": example["points"],
-        "boxes": example["boxes"],
-        "size": example["size"]
-    }
-    for example in examples
-}
+example_paths = [{"image": 'images/' + path} for path in os.listdir('images')]
 
 theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="emerald")
 with gr.Blocks(theme=theme, title="🔍 Compare SAM vs SAM-HQ") as demo:
@@ -113,53 +102,15 @@ with gr.Blocks(theme=theme, title="🔍 Compare SAM vs SAM-HQ") as demo:
     gr.Markdown("Compare the performance of SAM and SAM-HQ on various images. Click on an example to load it")
     gr.Markdown("[SAM-HQ](https://huggingface.co/syscv-community/sam-hq-vit-huge) - [SAM](https://huggingface.co/facebook/sam-vit-huge)")
 
-    with gr.Row():
-        image_input = gr.Image(
-            type="pil",
-            label="Example image (click below to load)",
-            interactive=False,
-            height=500,
-            show_label=True
-        )
-    
-    gr.Examples(
-        examples=example_paths,
-        inputs=[image_input],
-        label="Click an example to try 👇",
-    )
-
+    print('example_paths', example_paths)
     result_html = gr.HTML(elem_id="result-html")
-
-    with gr.Row():
-        points_input = gr.Dataframe(
-            headers=["x", "y"],
-            label="Points",
-            datatype=["number", "number"],
-            col_count=(2, "fixed")
-        )
-        boxes_input = gr.Dataframe(
-            headers=["x0", "y0", "x1", "y1"],
-            label="Boxes",
-            datatype=["number", "number", "number", "number"],
-            col_count=(4, "fixed")
-        )
-
-    def on_image_change(image):
-        for example_path, example_data in example_data_map.items():
-            print(image.size)
-            if example_data["size"] == list(image.size):
-                return example_data["points"], example_data["boxes"]
-        return [], []
-
-    image_input.change(
-        fn=on_image_change,
-        inputs=[image_input],
-        outputs=[points_input, boxes_input]
+    gr.Interface(
+        fn=process_inputs,
+        #examples=example_paths,
+        inputs=ImagePrompter(show_label=False),
+        outputs=result_html,
     )
 
-    compare_button = gr.Button("Compare points and masks")
-    compare_button.click(fn=compare_images_points_and_masks, inputs=[image_input, boxes_input, points_input], outputs=result_html)
-
     gr.HTML("""
         <style>
         #result-html {
@@ -174,4 +125,5 @@ with gr.Blocks(theme=theme, title="🔍 Compare SAM vs SAM-HQ") as demo:
         </style>
     """)
 
+
 demo.launch()

requirements.txt

@@ -5,4 +5,6 @@ pillow
 torch
 git+https://github.com/huggingface/transformers.git
 matplotlib
-numpy
+numpy
+gradio-image-prompter
+pydantic==2.10.6