Update yolo_model.py

yolo_model.py

@@ -12,7 +12,6 @@ _yolo_model: Optional[YOLO] = None  # cached model
 
 
 def get_yolo_model() -> YOLO:
-    """Load YOLO once and reuse."""
     global _yolo_model
     if _yolo_model is None:
         if not _WEIGHTS_PATH.exists():
@@ -22,7 +21,7 @@ def get_yolo_model() -> YOLO:
 
 
 def run_yolo_on_page(image: Image.Image, conf_threshold: float = 0.3) -> Dict[str, Any]:
-    """Run YOLO on full page and return annotated image + detections list."""
+    """Run YOLO on full page and return annotated image + detections + crops."""
     image = image.convert("RGB")
     img_rgb = np.array(image)
 
@@ -33,6 +32,7 @@ def run_yolo_on_page(image: Image.Image, conf_threshold: float = 0.3) -> Dict[st
     annotated_image = Image.fromarray(annotated)
 
     detections: List[Dict[str, Any]] = []
+    crops: List[Dict[str, Any]] = []
     names = results.names
 
     for box in results.boxes:
@@ -54,7 +54,22 @@ def run_yolo_on_page(image: Image.Image, conf_threshold: float = 0.3) -> Dict[st
             }
         )
 
+        if x2 > x1 and y2 > y1:
+            crop_np = img_rgb[y1:y2, x1:x2]
+            if crop_np.size > 0:
+                crop_img = Image.fromarray(crop_np)
+                crops.append(
+                    {
+                        "cls_id": cls_id,
+                        "cls_name": cls_name,
+                        "conf": conf,
+                        "box": [x1, y1, x2, y2],
+                        "crop_image": crop_img,
+                    }
+                )
+
     return {
         "annotated_image": annotated_image,
         "detections": detections,
+        "crops": crops,
     }