Update app.py

app.py

@@ -8,12 +8,12 @@ headers = {
     "x-api-key": os.getenv("JIGSAWSTACK_API_KEY")
 }
 
-def detect_objects(image_url=None, file_store_key=None):
+def detect_objects(image_url=None, file_store_key=None, prompts=None, features=None, annotated_image=False):
     if not image_url and not file_store_key:
-        return "❌ Please provide either an image URL or file store key.", [], "", ""
+        return "❌ Please provide either an image URL or file store key.", "", "", None
 
     if image_url and file_store_key:
-        return "❌ Provide only one: image URL or file store key.", [], "", ""
+        return "❌ Provide only one: image URL or file store key.", "", "", None
 
     try:
         payload = {}
@@ -21,34 +21,55 @@ def detect_objects(image_url=None, file_store_key=None):
             payload["url"] = image_url
         if file_store_key:
             payload["file_store_key"] = file_store_key
+        
+        # Add optional parameters
+        if prompts:
+            payload["prompts"] = prompts
+        if features:
+            payload["features"] = features
+        
+        # Always return annotated image
+        payload["annotated_image"] = True
+        
+        # Always use url as return_type
+        payload["return_type"] = "url"
 
         response = requests.post(f"{BASE_URL}/ai/object_detection", headers=headers, json=payload)
         if response.status_code != 200:
-            return f"❌ Error: {response.status_code} - {response.text}", [], "", ""
+            return f"❌ Error: {response.status_code} - {response.text}", "", "", None
 
         result = response.json()
         if not result.get("success"):
-            return "❌ Detection failed.", [], "", ""
+            return "❌ Detection failed.", "", "", None
 
         status = "✅ Detection successful!"
-        tags = result.get("tags", [])
         objects = result.get("objects", [])
-        description = f"Image Size: {result.get('width')} x {result.get('height')}\n\n"
+        annotated_image_url = result.get("annotated_image")
+        
+        # Create description with object details
+        description = f"Image Size: {result.get('width', 'Unknown')} x {result.get('height', 'Unknown')}\n\n"
+        description += f"Total Objects Detected: {len(objects)}\n\n"
 
-        for obj in objects:
+        for i, obj in enumerate(objects):
             bounds = obj.get("bounds", {})
+            label = obj.get("label", "Unknown")
             bound_text = ""
-            if bounds.get("top_left") and bounds.get("top_right"):
-                tl = bounds["top_left"]
-                tr = bounds["top_right"]
-                bound_text = f"Bounds: ({tl['x']}, {tl['y']}) to ({tr['x']}, {tr['y']})"
-            description += f"• {obj['name']} (Confidence: {obj['confidence']:.2f})\n  {bound_text}\n"
+            
+            if bounds:
+                width = bounds.get("width", "Unknown")
+                height = bounds.get("height", "Unknown")
+                top_left = bounds.get("top_left", {})
+                if top_left:
+                    x, y = top_left.get("x", "?"), top_left.get("y", "?")
+                    bound_text = f"Position: ({x}, {y}), Size: {width}x{height}"
+            
+            description += f"• {label}\n  {bound_text}\n"
 
         raw_json = json.dumps(result, indent=2)
-        return status, tags, description.strip(), raw_json
+        return status, description.strip(), raw_json, annotated_image_url
 
     except Exception as e:
-        return f"❌ Error: {str(e)}", [], "", ""
+        return f"❌ Error: {str(e)}", "", "", None
 
 with gr.Blocks() as demo:
     gr.Markdown("""
@@ -64,13 +85,22 @@ with gr.Blocks() as demo:
             input_type = gr.Radio(choices=["Image URL", "File Store Key"], value="Image URL", label="Input Type")
             image_url = gr.Textbox(label="Image URL", placeholder="https://example.com/image.jpg", visible=True)
             file_store_key = gr.Textbox(label="File Store Key", placeholder="my-image.jpg", visible=False)
+            
+            # Advanced options
+            prompts = gr.Textbox(label="Prompts (comma-separated)", placeholder="wine glass, bottle, cup", info="Targeted object detection prompts")
+            features = gr.CheckboxGroup(choices=["object_detection", "gui"], value=["object_detection"], label="Features")
+            annotated_image = gr.Checkbox(label="Return Annotated Image", value=True)
+            
             detect_btn = gr.Button("🔍 Detect Objects")
             clear_btn = gr.Button("Clear")
 
         with gr.Column():
             status_box = gr.Textbox(label="Status", interactive=False)
-            tag_display = gr.Label(label="Detected Tags")
             desc_display = gr.Textbox(label="Object Details", lines=10, interactive=False)
+            
+            # Annotated image display
+            annotated_image_display = gr.Image(label="Annotated Image", visible=False)
+            
             json_box = gr.Accordion("Raw JSON Response", open=False)
             with json_box:
                 json_output = gr.Textbox(show_label=False, lines=20, interactive=False)
@@ -83,20 +113,42 @@ with gr.Blocks() as demo:
 
     input_type.change(fn=toggle_inputs, inputs=input_type, outputs=[image_url, file_store_key])
 
-    def on_detect(input_mode, url, key):
+    def on_detect(input_mode, url, key, prompts_text, features_list, annotated):
+        # Parse prompts
+        prompts_list = None
+        if prompts_text.strip():
+            prompts_list = [p.strip() for p in prompts_text.split(",") if p.strip()]
+        
         if input_mode == "Image URL":
-            return detect_objects(image_url=url.strip())
+            return detect_objects(
+                image_url=url.strip(),
+                prompts=prompts_list,
+                features=features_list,
+                annotated_image=annotated
+            )
         else:
-            return detect_objects(file_store_key=key.strip())
+            return detect_objects(
+                file_store_key=key.strip(),
+                prompts=prompts_list,
+                features=features_list,
+                annotated_image=annotated
+            )
 
-    detect_btn.click(fn=on_detect, inputs=[input_type, image_url, file_store_key],
-                     outputs=[status_box, tag_display, desc_display, json_output])
+    def update_annotated_image_visibility(annotated):
+        return gr.update(visible=annotated)
+
+    detect_btn.click(fn=on_detect, inputs=[
+        input_type, image_url, file_store_key, prompts, features, annotated_image
+    ], outputs=[status_box, desc_display, json_output, annotated_image_display])
+
+    annotated_image.change(fn=update_annotated_image_visibility, inputs=annotated_image, outputs=annotated_image_display)
 
     def clear_all():
-        return "Image URL", "", "", "", "", ""
+        return "Image URL", "", "", "", "", ["object_detection"], False, "", "", "", None
 
     clear_btn.click(fn=clear_all, inputs=[], outputs=[
-        input_type, image_url, file_store_key, status_box, desc_display, json_output
+        input_type, image_url, file_store_key, prompts, features, annotated_image,
+        status_box, desc_display, json_output, annotated_image_display
     ])
 
 demo.launch()