Update app.py

app.py

@@ -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"] = "AIzaSyDOBd0_yNLckwsZJrpb9-CqTHFUx0Ah3R8"  # Replace with your actual API key
 gemini_model = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
 
 # Load YOLO model
@@ -31,6 +31,7 @@ def encode_image_to_base64(image):
     return base64.b64encode(img_buffer).decode('utf-8')
 
 def analyze_image_with_gemini(current_image):
+    """Send image to Gemini API for analysis."""
     if current_image is None:
         return "No image available for analysis."
     current_image_data = encode_image_to_base64(current_image)
@@ -47,15 +48,18 @@ def analyze_image_with_gemini(current_image):
         return f"Error processing image: {e}"
 
 def save_crop_image(crop, track_id):
+    """Save cropped image of detected bottle."""
     filename = f"{crop_folder}/{track_id}.jpg"
     cv2.imwrite(filename, crop)
     return filename
 
 def process_crop_image(crop, track_id):
+    """Process image asynchronously using Gemini AI."""
     response = analyze_image_with_gemini(crop)
     st.session_state["responses"].append((track_id, response))
 
 def process_video(uploaded_file):
+    """Process uploaded video, detect objects, and create an output video."""
     if not uploaded_file:
         return None
     
@@ -65,15 +69,22 @@ def process_video(uploaded_file):
         f.write(video_bytes)
     
     cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        st.error("Error: Could not open video file.")
+        return None
+
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     output_path = "output_video.mp4"
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(output_path, fourcc, 20.0, (1020, 500))
-    
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
     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()
@@ -84,22 +95,31 @@ def process_video(uploaded_file):
                     crop = frame[y1:y2, x1:x2]
                     save_crop_image(crop, track_id)
                     threading.Thread(target=process_crop_image, args=(crop, track_id)).start()
-                processed_track_ids.add(track_id)
+                    processed_track_ids.add(track_id)
+
         out.write(frame)
+
     cap.release()
     out.release()
+
     return output_path
 
+# Streamlit UI
 st.title("Bottle Label Checking using YOLO & Gemini AI")
-st.sidebar.header("Upload a video")
+st.sidebar.header("Upload a Video")
 uploaded_file = st.sidebar.file_uploader("Choose a video file", type=["mp4", "avi", "mov"])
+
 if "responses" not in st.session_state:
     st.session_state["responses"] = []
+
 if uploaded_file:
-    st.sidebar.write("Processing...")
+    st.sidebar.write("Processing video, please wait...")
     output_video_path = process_video(uploaded_file)
-    st.sidebar.success("Processing completed!")
-    st.video(output_video_path)
+
+    if output_video_path:
+        st.sidebar.success("Processing completed!")
+        st.video(output_video_path)
+    
     st.subheader("AI Analysis Results")
     for track_id, response in st.session_state["responses"]:
         st.write(f"**Track ID {track_id}:** {response}")