feat: YOLO image segmentation with Gradio UI

app.py

@@ -1,7 +1,63 @@
 import gradio as gr
+import cv2
+import numpy as np
+from ultralytics import YOLO
 
-def greet(name):
-    return "Hello " + name + "!!"
+def detect_and_crop_spines(image, model_path):
+    # Cache the YOLO model to avoid reloading it for every inference
+    if not hasattr(detect_and_crop_spines, 'model'):
+        detect_and_crop_spines.model = YOLO(model_path)
+    
+    # Convert gradio image to CV2 format if needed
+    if isinstance(image, str):
+        image = cv2.imread(image)
+    elif isinstance(image, np.ndarray) and len(image.shape) == 3 and image.shape[2] == 3:
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
+    # Add error handling for image loading
+    if image is None:
+        raise ValueError("Failed to load image")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+    # Use confidence threshold and add error handling for predictions
+    try:
+        results = detect_and_crop_spines.model.predict(image, show=False, conf=0.25)
+    except Exception as e:
+        return None, f"Error during detection: {str(e)}"
+
+    annotated_image = image.copy()
+    
+    for det in enumerate(results[0].boxes.xyxy.cpu().numpy()):
+        # Get corner points
+        points = np.array([[det[0], det[1]], [det[2], det[1]], 
+                          [det[2], det[3]], [det[0], det[3]]], dtype=np.float32)
+        
+        # Get minimum area rectangle
+        rect = cv2.minAreaRect(points)
+        box = cv2.boxPoints(rect)
+        box = np.int0(box)
+        
+        # Draw rotated rectangle
+        cv2.drawContours(annotated_image, [box], 0, (0, 255, 0), 2)
+    
+    # Convert back to RGB for Gradio
+    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+    
+    return annotated_image, f"Detected {len(results[0].boxes)} book spines"
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=detect_and_crop_spines,
+    inputs=[
+        gr.Image(),
+        gr.Text(label="YOLO Model Path", value="yolo11x.pt")
+    ],
+    outputs=[
+        gr.Image(label="Annotated Image"),
+        gr.Text(label="Detection Results")
+    ],
+    title="Book Spine Detector",
+    description="Upload an image of a bookshelf to detect book spines using YOLO"
+)
+
+# Launch the app
 demo.launch()

requirements.txt

@@ -4,9 +4,12 @@ anyio==4.8.0
 certifi==2024.12.14
 charset-normalizer==3.4.1
 click==8.1.8
+contourpy==1.3.1
+cycler==0.12.1
 fastapi==0.115.6
 ffmpy==0.5.0
 filelock==3.16.1
+fonttools==4.55.3
 fsspec==2024.12.0
 gradio==5.10.0
 gradio_client==1.5.3
@@ -16,18 +19,26 @@ httpx==0.28.1
 huggingface-hub==0.27.1
 idna==3.10
 Jinja2==3.1.5
+kiwisolver==1.4.8
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
+matplotlib==3.10.0
 mdurl==0.1.2
-numpy==2.2.1
+mpmath==1.3.0
+networkx==3.4.2
+numpy==1.26.4
+opencv-python==4.10.0.84
 orjson==3.10.13
 packaging==24.2
 pandas==2.2.3
 pillow==11.1.0
+psutil==6.1.1
+py-cpuinfo==9.0.0
 pydantic==2.10.4
 pydantic_core==2.27.2
 pydub==0.25.1
 Pygments==2.19.1
+pyparsing==3.2.1
 python-dateutil==2.9.0.post0
 python-multipart==0.0.20
 pytz==2024.2
@@ -36,16 +47,23 @@ requests==2.32.3
 rich==13.9.4
 ruff==0.8.6
 safehttpx==0.1.6
+scipy==1.15.0
+seaborn==0.13.2
 semantic-version==2.10.0
 shellingham==1.5.4
 six==1.17.0
 sniffio==1.3.1
 starlette==0.41.3
+sympy==1.13.1
 tomlkit==0.13.2
+torch==2.5.1
+torchvision==0.20.1
 tqdm==4.67.1
 typer==0.15.1
 typing_extensions==4.12.2
 tzdata==2024.2
+ultralytics==8.3.58
+ultralytics-thop==2.0.13
 urllib3==2.3.0
 uvicorn==0.34.0
 websockets==14.1