Create app.py

app.py

@@ -0,0 +1,136 @@
+import gradio as gr
+import cv2
+import numpy as np
+import os
+import time
+import threading
+import base64
+from ultralytics import YOLO
+from langchain_core.messages import HumanMessage
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+# Set up Google API Key
+os.environ["GOOGLE_API_KEY"] = "YOUR_GOOGLE_API_KEY"  # Replace with your API Key
+gemini_model = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
+
+# Load YOLO model
+yolo_model = YOLO("best.pt")
+names = yolo_model.names
+
+# Constants for ROI detection
+cx1 = 491
+offset = 8
+current_date = time.strftime("%Y-%m-%d")
+crop_folder = f"crop_{current_date}"
+if not os.path.exists(crop_folder):
+    os.makedirs(crop_folder)
+
+# Track processed IDs to avoid duplicate processing
+processed_track_ids = set()
+lock = threading.Lock()  # Ensure thread-safe operations
+
+def encode_image_to_base64(image):
+    _, img_buffer = cv2.imencode('.jpg', image)
+    return base64.b64encode(img_buffer).decode('utf-8')
+
+def analyze_image_with_gemini(current_image):
+    if current_image is None:
+        return "No image available for analysis."
+    
+    current_image_data = encode_image_to_base64(current_image)
+    message = HumanMessage(
+        content=[
+            {"type": "text", "text": "Analyze this image and check if the label is present on the bottle. Return results in a structured format."},
+            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{current_image_data}"}, "description": "Detected product"}
+        ]
+    )
+    try:
+        response = gemini_model.invoke([message])
+        return response.content
+    except Exception as e:
+        return f"Error processing image: {e}"
+
+def save_crop_image(crop, track_id):
+    filename = f"{crop_folder}/{track_id}.jpg"
+    cv2.imwrite(filename, crop)
+    return filename
+
+def process_crop_image(crop, track_id, responses):
+    response = analyze_image_with_gemini(crop)
+    responses.append((track_id, response))
+
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    output_path = "output_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, 20.0, (1020, 500))
+    
+    responses = []
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.resize(frame, (1020, 500))
+        
+        results = yolo_model.track(frame, persist=True)
+        if results[0].boxes is not None:
+            boxes = results[0].boxes.xyxy.int().cpu().tolist()
+            track_ids = results[0].boxes.id.int().cpu().tolist() if results[0].boxes.id is not None else [-1] * len(boxes)
+
+            for box, track_id in zip(boxes, track_ids):
+                with lock:  # Prevent race condition
+                    if track_id not in processed_track_ids:
+                        x1, y1, x2, y2 = box
+                        crop = frame[y1:y2, x1:x2]
+                        save_crop_image(crop, track_id)
+                        threading.Thread(target=process_crop_image, args=(crop, track_id, responses)).start()
+                        processed_track_ids.add(track_id)
+
+        out.write(frame)
+    
+    cap.release()
+    out.release()
+    return output_path, responses
+
+def process_and_return(video_file):
+    if not video_file:
+        return None, "No video uploaded."
+
+    video_path = "uploaded_video.mp4"
+    with open(video_path, "wb") as f:
+        f.write(video_file)
+
+    output_video_path, analysis_results = process_video(video_path)
+
+    results_text = "\n".join([f"**Track ID {track_id}:** {response}" for track_id, response in analysis_results])
+    
+    return output_video_path, results_text
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Bottle Label Checking using YOLO & Gemini AI")
+
+    with gr.Row():
+        video_input = gr.File(label="Upload a video", type="binary")
+        process_button = gr.Button("Process Video")
+
+    with gr.Row():
+        video_output = gr.Video(label="Processed Video")
+        download_button = gr.File(label="Download Processed Video")
+
+    analysis_results = gr.Markdown(label="AI Analysis Results")
+
+    process_button.click(
+        fn=process_and_return,
+        inputs=video_input,
+        outputs=[video_output, analysis_results]
+    )
+
+    download_button.change(
+        fn=lambda x: x if x else None,
+        inputs=video_output,
+        outputs=download_button
+    )
+
+demo.launch()