changes

app.py

@@ -20,13 +20,14 @@ CONFIDENCE_THRESHOLD = 0.3  # Confidence threshold for filtering predictions
 GRID_SIZE = (3, 3)  # 3x3 segmentation
 
 def detect_components(image):
-    """ Detect components in an uploaded image with three passes. """
+    """ Detect components in an uploaded image with three passes, removing detected areas after each pass. """
     
     original_image = image.convert("RGB")
     width, height = original_image.size
     seg_w, seg_h = width // GRID_SIZE[1], height // GRID_SIZE[0]
     
     def process_detection(image, pass_num):
+        """ Detect objects in an image segment and remove them if found. """
         final_image = image.copy()
         draw_final = ImageDraw.Draw(final_image)
         total_counts = defaultdict(int)
@@ -41,6 +42,7 @@ def detect_components(image):
                 segment_path = f"segment_{row}_{col}_pass{pass_num}.png"
                 segment.save(segment_path)
                 
+                # Run inference
                 result = CLIENT.infer(segment_path, model_id=MODEL_ID)
                 filtered_predictions = [pred for pred in result["predictions"] if pred["confidence"] >= CONFIDENCE_THRESHOLD]
                 
@@ -48,10 +50,13 @@ def detect_components(image):
                     sx, sy, sw, sh = obj["x"], obj["y"], obj["width"], obj["height"]
                     class_name = obj["class"]
                     total_counts[class_name] += 1
-                    
+
+                    # Convert segment coordinates to full image coordinates
                     x_min_full, y_min_full = x1 + sx - sw // 2, y1 + sy - sh // 2
                     x_max_full, y_max_full = x1 + sx + sw // 2, y1 + sy + sh // 2
                     detected_boxes.append((x_min_full, y_min_full, x_max_full, y_max_full))
+                    
+                    # Draw bounding box
                     draw_final.rectangle([x_min_full, y_min_full, x_max_full, y_max_full], outline="green", width=2)
         
         return final_image, total_counts, detected_boxes
@@ -59,7 +64,7 @@ def detect_components(image):
     # First pass detection
     image_after_pass1, counts_pass1, detected_boxes = process_detection(original_image, pass_num=1)
 
-    # Remove detected areas for second pass
+    # Mask detected areas for the second pass
     image_after_removal1 = original_image.copy()
     draw_removal1 = ImageDraw.Draw(image_after_removal1)
     for box in detected_boxes:
@@ -68,7 +73,7 @@ def detect_components(image):
     # Second pass detection
     image_after_pass2, counts_pass2, detected_boxes = process_detection(image_after_removal1, pass_num=2)
 
-    # Remove detected areas for third pass
+    # Mask detected areas for the third pass
     image_after_removal2 = image_after_removal1.copy()
     draw_removal2 = ImageDraw.Draw(image_after_removal2)
     for box in detected_boxes:
@@ -81,11 +86,8 @@ def detect_components(image):
     final_counts = defaultdict(int)
     for label in set(counts_pass1) | set(counts_pass2) | set(counts_pass3):
         final_counts[label] = counts_pass1.get(label, 0) + counts_pass2.get(label, 0) + counts_pass3.get(label, 0)
-    
-    # Ensure actual counts are displayed as numbers (no text formatting)
-    formatted_counts = {k: v for k, v in final_counts.items()}  
 
-    return image_after_pass1, image_after_pass2, image_after_pass3, formatted_counts
+    return image_after_pass1, image_after_pass2, image_after_pass3, dict(final_counts)
 
 # Gradio Interface
 interface = gr.Interface(