Update app.py

app.py

@@ -5,7 +5,7 @@ from ultralytics import YOLO
 from PIL import Image
 import pandas as pd
 
-# Model labels
+# Model labels for characters and license plates
 model1Labels = {0: 'single_number_plate', 1: 'double_number_plate'}
 
 model2Labels = {
@@ -18,6 +18,7 @@ model2Labels = {
 model = YOLO("models/LP-detection.pt")
 model2 = YOLO("models/Charcter-LP.pt")
 
+# Function to process license plate and detect text
 def prediction(image):
     result = model.predict(source=image, conf=0.5)
     boxes = result[0].boxes
@@ -85,11 +86,13 @@ def prediction(image):
 
     return lp_number, img_lp_final  
 
+# Function to process the video
 def process_video(video_file):
     # Open the video
     cap = cv2.VideoCapture(video_file)
+    
     if not cap.isOpened():  # Check if video was opened successfully
-        return None, "Error: Unable to open video."
+        return None, "Error: Unable to open video. Please check the file format or path."
 
     license_plate_texts = []
     processed_frames = []
@@ -100,12 +103,16 @@ def process_video(video_file):
         if not ret:
             break  # Exit if no frame is read
 
+        # Perform license plate detection and character recognition
         license_plate_text, cropped_plate_img = prediction(frame)
+
+        # Collect license plate text
         license_plate_texts.append(" ".join(license_plate_text))  # Join the list of texts into a single string
         processed_frames.append(cropped_plate_img)
 
-    if not processed_frames:  # If no frames were processed
-        return None, "Error: No frames processed. Check the video file."
+    # Check if frames were processed
+    if not processed_frames:
+        return None, "Error: No frames processed. Check the video file or format."
 
     # Save detected texts to Excel
     df = pd.DataFrame(license_plate_texts, columns=["License Plate"])
@@ -114,14 +121,22 @@ def process_video(video_file):
     # Save processed video with license plates highlighted
     output_video_path = 'processed_video.mp4'
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
-    out = cv2.VideoWriter(output_video_path, fourcc, 20.0, (frame.shape[1], frame.shape[0]))
 
-    for processed_frame in processed_frames:
-        out.write(processed_frame)
-    cap.release()
-    out.release()
+    # Ensure frame shape is valid
+    if len(processed_frames) > 0:
+        frame = processed_frames[0]
+        height, width, _ = frame.shape
+        out = cv2.VideoWriter(output_video_path, fourcc, 20.0, (width, height))
+
+        for processed_frame in processed_frames:
+            out.write(processed_frame)
+
+        cap.release()
+        out.release()
 
-    return output_video_path, "detected_license_plates.xlsx"
+        return output_video_path, "detected_license_plates.xlsx"
+    else:
+        return None, "Error: No valid frames found."
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -129,9 +144,9 @@ with gr.Blocks() as demo:
     gr.Markdown("Upload a video to get the license number of vehicles detected in each frame.")
 
     with gr.Row():
-        video_input = gr.File(label="Upload Video", type="filepath")  # Corrected the file type
-        video_output = gr.Video(label="Processed Video")
-        excel_output = gr.File(label="Excel File with Detected License Plates")
+        video_input = gr.Video(label="Upload Video", type="filepath")  # Use Video for file input
+        video_output = gr.Video(label="Processed Video")  # Output the processed video
+        excel_output = gr.File(label="Excel File with Detected License Plates")  # Output the Excel file with detected text
 
     video_input.upload(process_video, inputs=video_input, outputs=[video_output, excel_output])