Update app.py

app.py

@@ -6,14 +6,35 @@ from fpdf import FPDF
 import tempfile
 import cv2
 import os
+import torch.serialization
+from ultralytics.nn.tasks import DetectionModel
+import numpy as np
+
+# Allowlist the DetectionModel class to fix UnpicklingError
+torch.serialization.add_safe_globals([DetectionModel])
 
 # Load the pre-trained YOLOv8n model
-model = YOLO('yolov8n.pt')
+model = YOLO('model/yolov8n.pt')
 model.overrides['conf'] = 0.25  # Confidence threshold
 model.overrides['iou'] = 0.45   # IoU threshold
 model.overrides['agnostic_nms'] = False
 model.overrides['max_det'] = 1000
 
+def detect_empty_spots(img_width, img_height, bottle_bboxes, min_gap=50):
+    """Simple logic to infer empty spots based on gaps between bottles"""
+    empty_bboxes = []
+    # Assume shelf spans full image width; split into horizontal segments
+    x_coords = sorted([bbox[0] for bbox in bottle_bboxes] + [bbox[2] for bbox in bottle_bboxes])
+    x_coords = [0] + x_coords + [img_width]
+    
+    for i in range(len(x_coords) - 1):
+        gap_start = x_coords[i]
+        gap_end = x_coords[i + 1]
+        if gap_end - gap_start > min_gap:  # If gap is large enough, consider it an empty spot
+            empty_bboxes.append([gap_start, 0, gap_end, img_height])  # Full height for simplicity
+    
+    return empty_bboxes
+
 def process_input(input_file):
     # Handle video: Extract middle frame
     if input_file.lower().endswith(('.mp4', '.avi', '.mov')):
@@ -35,25 +56,30 @@ def process_input(input_file):
     boxes = results[0].boxes
     class_names = results[0].names
 
-    # Count bottles (assuming class 'bottle' from COCO dataset)
+    # Count bottles and collect their bounding boxes
     num_present = 0
-    num_empty = 0  # Note: yolov8n.pt doesn't detect empty shelves
-    empty_bboxes = []
+    bottle_bboxes = []
     for box in boxes:
         cls_id = int(box.cls.item())
         cls_name = class_names[cls_id]
         if cls_name == 'bottle':
             num_present += 1
-        # Empty spots not detected; you can add logic to infer empty spaces later
+            bottle_bboxes.append(box.xyxy[0].cpu().numpy())  # [x1, y1, x2, y2]
 
-    # Load image for annotation (highlight bottles instead of empty spots for now)
+    # Load image for annotation
     img = Image.open(process_img_path)
+    img_width, img_height = img.size
     draw = ImageDraw.Draw(img)
-    for box in boxes:
-        cls_id = int(box.cls.item())
-        if class_names[cls_id] == 'bottle':
-            bbox = box.xyxy[0].cpu().numpy()
-            draw.rectangle(((bbox[0], bbox[1]), (bbox[2], bbox[3])), outline="green", width=3)
+
+    # Draw bottle bounding boxes
+    for bbox in bottle_bboxes:
+        draw.rectangle(((bbox[0], bbox[1]), (bbox[2], bbox[3])), outline="green", width=3)
+
+    # Infer empty spots
+    empty_bboxes = detect_empty_spots(img_width, img_height, bottle_bboxes)
+    num_empty = len(empty_bboxes)
+    for bbox in empty_bboxes:
+        draw.rectangle(((bbox[0], bbox[1]), (bbox[2], bbox[3])), outline="red", width=3)
 
     # Generate PDF
     pdf = FPDF()
@@ -62,30 +88,45 @@ def process_input(input_file):
     pdf.cell(200, 10, txt="Wine Shop Inventory Report", ln=1, align='C')
     pdf.ln(5)
     pdf.cell(200, 10, txt=f"Number of bottles detected: {num_present}", ln=1)
-    pdf.cell(200, 10, txt="Note: Empty spot detection not supported in this model.", ln=1)
+    pdf.cell(200, 10, txt=f"Number of empty spots (inferred): {num_empty}", ln=1)
     pdf.ln(10)
 
     # Add annotated full image
     with tempfile.NamedTemporaryFile(suffix='.png') as tmp_annotated:
         img.save(tmp_annotated.name)
-        pdf.image(tmp_annotated.name, x=10, y=pdf.get_y(), w=180)
+        pdf.image(tmp_annotated.name, x=10, y=pdf.get_y(), w=180)  # Full width for readability
+
+    # Add screenshots of bottles and empty spots
     pdf.add_page()
-    pdf.cell(200, 10, txt="Detected Bottles", ln=1, align='C')
+    pdf.cell(200, 10, txt="Detected Bottles and Empty Spots", ln=1, align='C')
     pdf.ln(5)
     y_pos = pdf.get_y()
-    for i, box in enumerate(boxes):
-        if class_names[int(box.cls.item())] == 'bottle':
-            bbox = box.xyxy[0].cpu().numpy()
-            cropped_img = img.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
-            with tempfile.NamedTemporaryFile(suffix='.png') as tmp_crop:
-                cropped_img.save(tmp_crop.name)
-                pdf.image(tmp_crop.name, x=10, y=y_pos, w=90)
-                y_pos += 100
-                pdf.cell(200, 10, txt=f"Bottle {i+1} (Location: x1={int(bbox[0])}, y1={int(bbox[1])}, x2={int(bbox[2])}, y2={int(bbox[3])})", ln=1)
-                pdf.ln(5)
-                if y_pos > 200:
-                    pdf.add_page()
-                    y_pos = 10
+
+    # Bottles
+    for i, bbox in enumerate(bottle_bboxes):
+        cropped_img = img.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+        with tempfile.NamedTemporaryFile(suffix='.png') as tmp_crop:
+            cropped_img.save(tmp_crop.name)
+            pdf.image(tmp_crop.name, x=10, y=y_pos, w=90)  # Medium size for clear view
+            y_pos += 100
+            pdf.cell(200, 10, txt=f"Bottle {i+1} (Location: x1={int(bbox[0])}, y1={int(bbox[1])}, x2={int(bbox[2])}, y2={int(bbox[3])})", ln=1)
+            pdf.ln(5)
+            if y_pos > 200:
+                pdf.add_page()
+                y_pos = 10
+
+    # Empty spots
+    for i, bbox in enumerate(empty_bboxes):
+        cropped_img = img.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+        with tempfile.NamedTemporaryFile(suffix='.png') as tmp_crop:
+            cropped_img.save(tmp_crop.name)
+            pdf.image(tmp_crop.name, x=10, y=y_pos, w=90)
+            y_pos += 100
+            pdf.cell(200, 10, txt=f"Empty Spot {i+1} (Location: x1={int(bbox[0])}, y1={int(bbox[1])}, x2={int(bbox[2])}, y2={int(bbox[3])})", ln=1)
+            pdf.ln(5)
+            if y_pos > 200:
+                pdf.add_page()
+                y_pos = 10
 
     # Output PDF
     pdf_bytes = io.BytesIO()