Update app.py

app.py

@@ -328,6 +328,7 @@
 
 # demo.launch(share=True)
 
+# imports
 # imports
 import os
 import json
@@ -395,25 +396,54 @@ def detect_objects(query_text):
     image_path = save_temp_image(state.current_image)
     
     try:
-        # Use VisionAgent to detect objects
-        image = T.load_image(image_path)
-        
         # Clean query text to get the object name
         object_name = query_text[0].replace("a photo of ", "").strip()
         
-        # Detect objects using CountGD
-        detections = T.countgd_object_detection(object_name, image)
+        # Let VisionAgent handle the detection with its agent-based approach
+        # Create agent message for object detection
+        agent_message = [
+            AgentMessage(
+                role="user",
+                content=f"Count the number of {object_name} in this image. Only show detections with high confidence (>0.75).",
+                media=[image_path]
+            )
+        ]
+        
+        # Generate code using VisionAgent
+        code_context = agent.generate_code(agent_message)
         
-        # Visualize results
-        result_image = T.overlay_bounding_boxes(image, detections)
+        # Load the image for visualization
+        image = T.load_image(image_path)
+        
+        # Use multiple models for detection and get high confidence results
+        # First try the specialized detector
+        detections = T.countgd_object_detection(object_name, image, conf_threshold=0.75)
+        
+        # If no high-confidence detections, try the more general object detector
+        if not detections:
+            # Try a different model with the same high threshold
+            try:
+                detections = T.grounding_dino_detection(object_name, image, box_threshold=0.75)
+            except:
+                pass
+        
+        # Only keep high confidence detections
+        high_conf_detections = [det for det in detections if det.get("score", 0) > 0.75]
+        
+        # Visualize only high confidence results with clear labeling
+        result_image = T.overlay_bounding_boxes(
+            image, 
+            high_conf_detections,
+            labels=[f"{object_name}: {det['score']:.2f}" for det in high_conf_detections]
+        )
         
         # Convert result back to numpy array for display
         state.last_prediction = np.array(result_image)
         
         return {
-            "count": len(detections),
-            "confidence": [det["score"] for det in detections],
-            "message": f"Detected {len(detections)} {object_name}(s)"
+            "count": len(high_conf_detections),
+            "confidence": [det["score"] for det in high_conf_detections],
+            "message": f"Detected {len(high_conf_detections)} {object_name}(s) with high confidence (>0.75)"
         }
     except Exception as e:
         print(f"Error in detect_objects: {str(e)}")
@@ -511,8 +541,9 @@ def chat(message, image, history):
     # Extract objects to detect from user message
     objects_to_detect = message.lower()
     
-    # Format query for object detection
-    query = ["a photo of " + objects_to_detect.replace("count", "").replace("detect", "").replace("show", "").strip()]
+    # Format query for object detection - keep it simple and direct
+    cleaned_query = objects_to_detect.replace("count", "").replace("detect", "").replace("show", "").strip()
+    query = ["a photo of " + cleaned_query]
     
     messages.append({
         "role": "user",