Update app.py

app.py

@@ -42,30 +42,34 @@ def chat_with_gemini(message):
 def analyze_image(image_path):
     global chat_history
     try:
-        # Open the image file
+        # Open and preprocess the image
         image = Image.open(image_path).convert("RGB")
-        # Preprocess the image for DETR
         inputs = processor(images=image, return_tensors="pt")
 
         # Perform inference
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Get predictions
+        # Set a threshold for filtering low-confidence detections (e.g., 0.8)
+        threshold = 0.8
         target_sizes = torch.tensor([image.size[::-1]])  # (height, width)
-        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[0]
 
-        # If any objects detected, display labels
+        # Filter detections by confidence threshold
         if len(results["labels"]) > 0:
-            # Decode the labels to human-readable objects
             detected_objects = []
-            for label in results["labels"]:
-                # Map the label to the COCO dataset categories (Detr uses COCO labels by default)
-                detected_objects.append(str(label.item()))  # Get the integer value (COCO class id)
-            bot_response = f"Objects detected: {', '.join(detected_objects)}."
+            for idx, label in enumerate(results["labels"]):
+                if results["scores"][idx] >= threshold:  # Check if confidence score is high enough
+                    # Get the object label based on label index
+                    object_name = COCO_CLASSES[label.item()]  # Assuming COCO_CLASSES is available
+                    detected_objects.append(object_name)
+            if detected_objects:
+                bot_response = f"Objects detected: {', '.join(detected_objects)}."
+            else:
+                bot_response = "No objects detected above the confidence threshold."
         else:
             bot_response = "No objects detected."
-        
+
         chat_history.append(("Uploaded an image for analysis", bot_response))
         return chat_history
     except Exception as e:
@@ -74,6 +78,7 @@ def analyze_image(image_path):
         return chat_history
 
 
+
 # Build the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Ken Chatbot")