Update app.py

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 import cv2
 import numpy as np
 import os
@@ -10,7 +10,7 @@ from langchain_core.messages import HumanMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
 
 # Set up Google API Key
-os.environ["GOOGLE_API_KEY"] = "AIzaSyDOBd0_yNLckwsZJrpb9-CqTHFUx0Ah3R8"  # Replace with your 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
@@ -29,10 +29,6 @@ if not os.path.exists(crop_folder):
 processed_track_ids = set()
 lock = threading.Lock()  # Ensure thread-safe operations
 
-# Ensure session state is initialized
-if "responses" not in st.session_state:
-    st.session_state["responses"] = []
-
 def encode_image_to_base64(image):
     _, img_buffer = cv2.imencode('.jpg', image)
     return base64.b64encode(img_buffer).decode('utf-8')
@@ -59,28 +55,18 @@ def save_crop_image(crop, track_id):
     cv2.imwrite(filename, crop)
     return filename
 
-def process_crop_image(crop, track_id):
+def process_crop_image(crop, track_id, responses):
     response = analyze_image_with_gemini(crop)
-    
-    # Thread-safe way to update session state
-    if "responses" not in st.session_state:
-        st.session_state["responses"] = []
-    st.session_state["responses"] = st.session_state["responses"] + [(track_id, response)]
-
-def process_video(uploaded_file):
-    if not uploaded_file:
-        return None
-    
-    video_bytes = uploaded_file.read()
-    video_path = "uploaded_video.mp4"
-    with open(video_path, "wb") as f:
-        f.write(video_bytes)
-    
+    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:
@@ -98,37 +84,53 @@ def process_video(uploaded_file):
                         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)).start()
+                        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
+    return output_path, responses
+
+def process_and_return(video_file):
+    if not video_file:
+        return None, "No video uploaded."
 
-def get_video_download_link(file_path):
-    """Generates a download link for the processed video."""
-    with open(file_path, "rb") as file:
-        video_bytes = file.read()
-    b64 = base64.b64encode(video_bytes).decode()
-    return f'<a href="data:video/mp4;base64,{b64}" download="output_video.mp4">Download Processed Video</a>'
+    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
 
-st.title("Bottle Label Checking using YOLO & Gemini AI")
-st.sidebar.header("Upload a video")
-uploaded_file = st.sidebar.file_uploader("Choose a video file", type=["mp4", "avi", "mov"])
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Bottle Label Checking using YOLO & Gemini AI")
 
-if uploaded_file:
-    st.sidebar.write("Processing... Please wait!")
-    output_video_path = process_video(uploaded_file)
-    st.sidebar.success("Processing completed!")
+    with gr.Row():
+        video_input = gr.File(label="Upload a video", type="binary")
+        process_button = gr.Button("Process Video")
 
-    # Show processed video
-    st.video(output_video_path)
+    with gr.Row():
+        video_output = gr.Video(label="Processed Video")
+        download_button = gr.File(label="Download Processed Video")
 
-    # Provide download link for the video
-    st.markdown(get_video_download_link(output_video_path), unsafe_allow_html=True)
+    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
+    )
 
-    st.subheader("AI Analysis Results")
-    for track_id, response in st.session_state["responses"]:
-        st.write(f"**Track ID {track_id}:** {response}")
+demo.launch()