Update app.py

app.py

@@ -8,10 +8,11 @@ import cv2
 import os
 import torch.serialization
 from ultralytics.nn.tasks import DetectionModel
+from torch.nn.modules.container import Sequential
 import numpy as np
 
-# Allowlist the DetectionModel class to fix UnpicklingError
-torch.serialization.add_safe_globals([DetectionModel])
+# Allowlist required classes to fix UnpicklingError
+torch.serialization.add_safe_globals([DetectionModel, Sequential])
 
 # Load the pre-trained YOLOv8n model
 model = YOLO('model/yolov8n.pt')
@@ -21,22 +22,19 @@ 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
+        if gap_end - gap_start > min_gap:
+            empty_bboxes.append([gap_start, 0, gap_end, img_height])
     
     return empty_bboxes
 
 def process_input(input_file):
-    # Handle video: Extract middle frame
     if input_file.lower().endswith(('.mp4', '.avi', '.mov')):
         cap = cv2.VideoCapture(input_file)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -51,12 +49,10 @@ def process_input(input_file):
     else:
         process_img_path = input_file
 
-    # Run model inference
     results = model(process_img_path)
     boxes = results[0].boxes
     class_names = results[0].names
 
-    # Count bottles and collect their bounding boxes
     num_present = 0
     bottle_bboxes = []
     for box in boxes:
@@ -64,24 +60,20 @@ def process_input(input_file):
         cls_name = class_names[cls_id]
         if cls_name == 'bottle':
             num_present += 1
-            bottle_bboxes.append(box.xyxy[0].cpu().numpy())  # [x1, y1, x2, y2]
+            bottle_bboxes.append(box.xyxy[0].cpu().numpy())
 
-    # Load image for annotation
     img = Image.open(process_img_path)
     img_width, img_height = img.size
     draw = ImageDraw.Draw(img)
 
-    # 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()
     pdf.add_page()
     pdf.set_font("Arial", size=12)
@@ -91,23 +83,20 @@ def process_input(input_file):
     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)  # Full width for readability
+        pdf.image(tmp_annotated.name, x=10, y=pdf.get_y(), w=180)
 
-    # Add screenshots of bottles and empty spots
     pdf.add_page()
     pdf.cell(200, 10, txt="Detected Bottles and Empty Spots", ln=1, align='C')
     pdf.ln(5)
     y_pos = pdf.get_y()
 
-    # 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
+            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)
@@ -115,7 +104,6 @@ def process_input(input_file):
                 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:
@@ -128,18 +116,15 @@ def process_input(input_file):
                 pdf.add_page()
                 y_pos = 10
 
-    # Output PDF
     pdf_bytes = io.BytesIO()
     pdf.output(pdf_bytes)
     pdf_bytes.seek(0)
 
-    # Clean up temp files
     if 'temp_frame_path' in locals():
         os.remove(temp_frame_path)
 
     return pdf_bytes.getvalue()
 
-# Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Wine Shop CCTV Analyzer")
     gr.Markdown("Upload a CCTV image or video to analyze stock and generate a PDF report.")