Create app.py

app.py

@@ -0,0 +1,54 @@
+import cv2
+import torch
+from ultralytics import YOLO
+import gradio as gr
+
+# Load the pre-trained YOLO model (assuming 'best.pt' is a YOLOv5 model)
+model = YOLO("./data/best.pt")
+
+# Function to process video frames and count wine bottles
+def process_frame(frame):
+    # Perform inference on the frame
+    results = model(frame)
+    
+    # Extract results
+    detections = results.pandas().xywh[results.pandas().xywh['class'] == 0]  # Assuming '0' is the class for wine bottles
+    
+    # Count the number of wine bottles detected
+    bottle_count = len(detections)
+    return bottle_count
+
+# Classify stock based on bottle count
+def classify_stock(bottle_count):
+    if bottle_count > 50:
+        return "Full"
+    elif 20 <= bottle_count <= 50:
+        return "Medium"
+    else:
+        return "Low"
+
+# Video processing function to classify each frame and track stock level
+def classify_video(video):
+    cap = cv2.VideoCapture(video.name)
+    stock_status = None
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        bottle_count = process_frame(frame)
+        stock_status = classify_stock(bottle_count)
+    
+    cap.release()
+    return stock_status
+
+# Gradio interface to upload a video and classify stock
+def main(video_input):
+    return classify_video(video_input)
+
+# Creating the Gradio interface
+iface = gr.Interface(fn=main, inputs=gr.Video(), outputs="text")
+
+if __name__ == "__main__":
+    iface.launch()