Update app.py

app.py

@@ -23,9 +23,9 @@ def load_easyocr_reader():
 yolo_model = load_yolo_model()
 ocr_reader = load_easyocr_reader()
 
-# Function to process the uploaded image
+# Function to process image and detect license plates
 def process_image(image, confidence_threshold=0.5):
-    # Perform license plate detection using YOLO
+    # Perform license plate detection
     results = yolo_model(image, conf=confidence_threshold)
     annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
     st.image(annotated_image, caption="Detected License Plate(s)", use_container_width=True)
@@ -46,59 +46,83 @@ def process_image(image, confidence_threshold=0.5):
             text_results = ocr_reader.readtext(cropped_image_rgb, detail=0)
             detected_text = " ".join(text_results)
             st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
+            st.write(f"**Confidence Score:** {result.boxes.conf.cpu().numpy()[i]:.2f}")
 
 # Function to process video and detect license plates
-def process_video(video_path, confidence_threshold=0.5):
+def process_video(video_path, confidence_threshold=0.5, output_path="output_video.mp4"):
+    # Open the video file
     cap = cv2.VideoCapture(video_path)
-    
+
+    # Get video frame dimensions
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Create VideoWriter object to save the output video
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for mp4
+    out = cv2.VideoWriter(output_path, fourcc, 20.0, (frame_width, frame_height))  # 20 FPS
+
     if not cap.isOpened():
         st.error("Error opening video stream or file")
         return
-    
+
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
+
         results = yolo_model(frame, conf=confidence_threshold)
         annotated_frame = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
 
-        st.image(annotated_frame, caption="Processed Video Frame", use_container_width=True)
-
+        # Loop through detections and perform OCR
         for result in results:
             boxes = result.boxes.xyxy.cpu().numpy().astype(int)
             for i, box in enumerate(boxes):
                 x1, y1, x2, y2 = box
                 cropped_plate = frame[y1:y2, x1:x2]
                 cropped_rgb = cv2.cvtColor(cropped_plate, cv2.COLOR_BGR2RGB)
-                
+
                 # Perform OCR on the cropped image
                 text_results = ocr_reader.readtext(cropped_rgb, detail=0)
                 detected_text = " ".join(text_results)
+
+                # Optionally add detected text on the annotated frame
                 cv2.putText(annotated_frame, detected_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
 
-        st.image(annotated_frame, caption="Video Frame with OCR", use_container_width=True)
+        # Write the annotated frame to the output video
+        out.write(cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR))
 
     cap.release()
+    out.release()
+
+    st.success(f"Video processing complete. Output video saved to {output_path}")
+
+    # Provide a download link for the processed video
+    with open(output_path, "rb") as f:
+        st.download_button(label="Download Processed Video", data=f, file_name=output_path)
 
 # Sidebar input for file upload
-file = st.file_uploader("Upload an Image or Video", type=["jpg", "jpeg", "png", "mp4", "avi", "mov"])
+uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])
 
-if file is not None:
-    # Check file type
-    file_type = file.type
+if uploaded_file is not None:
+    # Check if it's an image or video
+    file_type = uploaded_file.type
 
-    # Process image
-    if file_type in ["image/jpeg", "image/png", "image/jpg"]:
-        image = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
-        process_image(image)
+    if file_type.startswith("image"):
+        # Read and process the image
+        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
+        confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
+        process_image(image, confidence_threshold)
 
-    # Process video
-    elif file_type in ["video/mp4", "video/avi", "video/mov"]:
-        video_bytes = file.read()
-        with open("/tmp/uploaded_video.mp4", "wb") as f:
+    elif file_type.startswith("video"):
+        # Save the uploaded video to a temporary file
+        video_bytes = uploaded_file.read()
+        video_path = "/tmp/uploaded_video.mp4"
+        with open(video_path, "wb") as f:
             f.write(video_bytes)
-        process_video("/tmp/uploaded_video.mp4")
+
+        # Process the video and save the output
+        output_path = "/tmp/output_video.mp4"
+        process_video(video_path, confidence_threshold=0.5, output_path=output_path)
 
 st.markdown("---")
 st.info("**Note:** This application uses EasyOCR for text recognition. Results may vary depending on image/video quality and lighting conditions.")