Update app.py

app.py

@@ -25,6 +25,9 @@ except Exception as e:
     raise RuntimeError(f"Error initializing TrOCR model: {e}")
 
 def detect_license_plate(frame):
+    """
+    Detect license plates in a video frame using DETR.
+    """
     # Convert the frame to a PIL image
     pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
     
@@ -32,7 +35,7 @@ def detect_license_plate(frame):
     inputs = detr_processor(images=pil_image, return_tensors="pt")
     outputs = detr_model(**inputs)
     
-    # Get detected objects and filter for license plates (e.g., class_id for "license plate")
+    # Get detected objects and filter for license plates
     logits = outputs.logits
     boxes = outputs.pred_boxes
     probas = logits.softmax(-1)[0, :, :-1]
@@ -47,6 +50,9 @@ def detect_license_plate(frame):
     return detected_boxes
 
 def recognize_text(plate_image):
+    """
+    Recognize text from a license plate image using TrOCR.
+    """
     # Convert the license plate image to a PIL image
     pil_image = Image.fromarray(cv2.cvtColor(plate_image, cv2.COLOR_BGR2RGB))
     
@@ -57,17 +63,23 @@ def recognize_text(plate_image):
     
     return text.strip()
 
-def main():
-    # Open video capture (0 for webcam or provide video file path)
-    cap = cv2.VideoCapture('road_video.mp4')  # Replace with your video file or 0 for webcam
-    
+def process_video(video_path, frame_skip=5):
+    """
+    Process a video to detect license plates and log entry/exit times.
+    """
+    cap = cv2.VideoCapture(video_path)
     vehicle_data = {}
-    
+    frame_count = 0
+
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
+
+        frame_count += 1
+        if frame_count % frame_skip != 0:
+            continue  # Skip frames to optimize processing time
+
         # Detect license plates
         detected_boxes = detect_license_plate(frame)
         
@@ -93,13 +105,13 @@ def main():
                 cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
                 cv2.putText(frame, license_plate, (x_min, y_min-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
         
-        # Display the frame
+        # Display the frame (optional, can be removed for headless environments)
         cv2.imshow('Vehicle Detection', frame)
         
         # Break on 'q' key press
         if cv2.waitKey(1) & 0xFF == ord('q'):
             break
-    
+
     cap.release()
     cv2.destroyAllWindows()
     
@@ -109,4 +121,5 @@ def main():
         print(f"License Plate: {plate}, Entry Time: {times['entry_time']}, Exit Time: {times['exit_time']}")
 
 if __name__ == "__main__":
-    main()
+    # Replace 'road_video.mp4' with the path to your video file or use 0 for webcam
+    process_video("road_video.mp4", frame_skip=5)