Update app.py

app.py

@@ -86,7 +86,7 @@ def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]
 
 def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if len(contours) is 0:
+    if len(contours) == 0:
         return np.array([])
     largest_contour = max(contours, key=cv2.contourArea)
     return largest_contour
@@ -190,11 +190,11 @@ def detections_to_json(detections):
         detections_list.append(detection_dict)
     return detections_list
 
-def crop_bounding_boxes(image: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
+def crop_bounding_boxes_with_yellow_background(image: np.ndarray, yellow_background: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
     crops = []
     for detection in detections:
         xmin, ymin, xmax, ymax = detection.box.xyxy
-        crop = image[ymin:ymax, xmin:xmax]
+        crop = yellow_background[ymin:ymax, xmin:xmax]
         crops.append(crop)
     return crops
 
@@ -206,7 +206,7 @@ def process_image(image, include_json, include_bboxes):
 
     results = [annotated_image]
     if include_bboxes:
-        crops = crop_bounding_boxes(np.array(original_image), detections)
+        crops = crop_bounding_boxes_with_yellow_background(np.array(original_image), yellow_background_with_insects, detections)
         results.extend(crops)
 
     if include_json:
@@ -215,6 +215,8 @@ def process_image(image, include_json, include_bboxes):
         with open(json_output_path, 'w') as json_file:
             json.dump(detections_json, json_file, indent=4)
         results.append(json.dumps(detections_json, separators=(',', ':')))
+    elif not include_bboxes:
+        results.append(None)
     
     return tuple(results)
 
@@ -235,12 +237,18 @@ with gr.Blocks() as demo:
 
     def update_outputs(image, include_json, include_bboxes):
         results = process_image(image, include_json, include_bboxes)
-        if include_bboxes:
+        if include_bboxes and include_json:
             annotated_img, *crops, json_txt = results
             return (annotated_img, json_txt, crops)
-        else:
+        elif include_bboxes:
+            annotated_img, *crops = results
+            return (annotated_img, None, crops)
+        elif include_json:
             annotated_img, json_txt = results
             return (annotated_img, json_txt, [])
+        else:
+            annotated_img, = results
+            return (annotated_img, None, [])
 
     submit_button.click(update_outputs, [image_input, include_json, include_bboxes], [annotated_output, json_output, crops_output])