Update app.py

app.py

@@ -81,47 +81,76 @@ def upload_to_s3(file_name, bucket, object_name=None):
     except NoCredentialsError:
         return False
 
+def crop_image(image_path, click_xy, crop_factor=0.5):
+    """Crop the image around the click point."""
+    image = Image.open(image_path)
+    width, height = image.size
+    crop_width, crop_height = int(width * crop_factor), int(height * crop_factor)
+
+    center_x, center_y = int(click_xy[0] * width), int(click_xy[1] * height)
+    left = max(center_x - crop_width // 2, 0)
+    upper = max(center_y - crop_height // 2, 0)
+    right = min(center_x + crop_width // 2, width)
+    lower = min(center_y + crop_height // 2, height)
+
+    cropped_image = image.crop((left, upper, right, lower))
+    cropped_image_path = f"cropped_{os.path.basename(image_path)}"
+    cropped_image.save(cropped_image_path)
+
+    return cropped_image_path
+
 @spaces.GPU
-def run_showui(image, query, session_id):
-    """Main function for inference."""
+def run_showui(image, query, session_id, iterations=2):
+    """Main function for iterative inference."""
     image_path = array_to_image_path(image, session_id)
     
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": _SYSTEM},
-                {"type": "image", "image": image_path, "min_pixels": MIN_PIXELS, "max_pixels": MAX_PIXELS},
-                {"type": "text", "text": query}
-            ],
-        }
-    ]
-
-    global model
-    model = model.to("cuda")
-    
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt"
-    )
-    inputs = inputs.to("cuda")
-
-    generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )[0]
-
-    click_xy = ast.literal_eval(output_text)
-    result_image = draw_point(image_path, click_xy, radius=10)
-    return result_image, str(click_xy), image_path
+    click_xy = None
+    images_during_iterations = []  # List to store images at each step
+
+    for _ in range(iterations):
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": _SYSTEM},
+                    {"type": "image", "image": image_path, "min_pixels": MIN_PIXELS, "max_pixels": MAX_PIXELS},
+                    {"type": "text", "text": query}
+                ],
+            }
+        ]
+
+        global model
+        model = model.to("cuda")
+
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt"
+        )
+        inputs = inputs.to("cuda")
+
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+
+        click_xy = ast.literal_eval(output_text)
+
+        # Draw point on the current image
+        result_image = draw_point(image_path, click_xy, radius=10)
+        images_during_iterations.append(result_image)  # Store the current image
+
+        # Crop the image for the next iteration
+        image_path = crop_image(image_path, click_xy)
+
+    return images_during_iterations, str(click_xy)
 
 def save_and_upload_data(image_path, query, session_id, is_example_image, votes=None):
     """Save the data to a JSON file and upload to S3."""
@@ -221,6 +250,10 @@ def build_demo(embed_mode, concurrency_count=1):
                 
                 Then upload/paste from clipboard 🤗
                 """)
+                
+                # Add a slider for iteration count
+                iteration_slider = gr.Slider(minimum=1, maximum=3, step=1, value=1, label="Refinement Steps")
+
                 textbox = gr.Textbox(
                     show_label=True,
                     placeholder="Enter a query (e.g., 'Click Nahant')",
@@ -258,13 +291,9 @@ def build_demo(embed_mode, concurrency_count=1):
                 )
 
             with gr.Column(scale=8):
-                output_img = gr.Image(type="pil", label="Output Image")
-                gr.HTML(
-                    """
-                    <p><strong>Note:</strong> The <span style="color: red;">red point</span> on the output image represents the predicted clickable coordinates.</p>
-                    """
-                )
-                output_coords = gr.Textbox(label="Clickable Coordinates")
+                # output_gallery = gr.Gallery(label="Iterative Refinement", object_fit="contain")
+                output_gallery = gr.Gallery(label="Iterative Refinement")
+                output_coords = gr.Textbox(label="Final Clickable Coordinates")
 
                 gr.HTML(
                     """
@@ -276,28 +305,28 @@ def build_demo(embed_mode, concurrency_count=1):
                     downvote_btn = gr.Button(value="👎 Too bad!", variant="secondary")
                     clear_btn = gr.Button(value="🗑️ Clear", interactive=True)
 
-            def on_submit(image, query, is_example_image):
+            def on_submit(image, query, iterations, is_example_image):
                 if image is None:
                     raise ValueError("No image provided. Please upload an image before submitting.")
                 
                 session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
                 
-                result_image, click_coords, image_path = run_showui(image, query, session_id)
+                images_during_iterations, click_coords = run_showui(image, query, session_id, iterations)
                 
-                save_and_upload_data(image_path, query, session_id, is_example_image)
+                save_and_upload_data(images_during_iterations[-1], query, session_id, is_example_image)
                 
-                return result_image, click_coords, image_path, session_id
+                return images_during_iterations, click_coords, session_id
 
             submit_btn.click(
                 on_submit,
-                [imagebox, textbox, is_example_dropdown],
-                [output_img, output_coords, state_image_path, state_session_id],
+                [imagebox, textbox, iteration_slider, is_example_dropdown],
+                [output_gallery, output_coords, state_session_id],
             )
 
             clear_btn.click(
-                lambda: (None, None, None, None, None, None),
+                lambda: (None, None, None, None),
                 inputs=None,
-                outputs=[imagebox, textbox, output_img, output_coords, state_image_path, state_session_id],
+                outputs=[imagebox, textbox, output_gallery, output_coords, state_session_id],
                 queue=False
             )
 
@@ -324,4 +353,4 @@ if __name__ == "__main__":
         server_port=7860,
         ssr_mode=False,
         debug=True,
-    )
+    )