Update app.py

app.py

@@ -4,133 +4,131 @@ import numpy as np
 import onnxruntime as ort
 from PIL import Image
 import tempfile
+import torch
+from ultralytics import YOLO
 
-# Define class labels
-CLASSES = {
-    0: "Vehicle",
-    1: "License_Plate"
-}
-
-# Load the ONNX model
+# Load models
 @st.cache_resource
-def load_model():
-    return ort.InferenceSession("model.onnx")
-
-ort_session = load_model()
-
-def preprocess_image(image, target_size=(640, 640)):
-    if isinstance(image, Image.Image):
-        image = np.array(image)
-    
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    original_shape = image.shape[:2]
-    image = cv2.resize(image, target_size)
-    image = image.astype(np.float32) / 255.0
-    image = np.transpose(image, (2, 0, 1))
-    image = np.expand_dims(image, axis=0)
-    return image, original_shape
-
-def postprocess_results(output, original_shape, confidence_threshold=0.25, iou_threshold=0.45):
-    if isinstance(output, (list, tuple)):
-        predictions = output[0]
-    elif isinstance(output, np.ndarray):
-        predictions = output
-    else:
-        raise ValueError(f"Unexpected output type: {type(output)}")
+def load_models():
+    license_plate_detector = YOLO('license_plate_detector.pt')
+    vehicle_detector = YOLO('yolov8n.pt')
+    ort_session = ort.InferenceSession("model.onnx")
+    return license_plate_detector, vehicle_detector, ort_session
 
-    if len(predictions.shape) == 4:
-        predictions = predictions.squeeze((0, 1))
-    elif len(predictions.shape) == 3:
-        predictions = predictions.squeeze(0)
+def draw_border(img, top_left, bottom_right, color=(0, 255, 0), thickness=10, line_length_x=200, line_length_y=200):
+    x1, y1 = top_left
+    x2, y2 = bottom_right
 
-    # Extract boxes, scores, and class_ids
-    boxes = predictions[:, :4]
-    scores = predictions[:, 4]
-    class_ids = predictions[:, 5]
+    # Draw corner lines
+    cv2.line(img, (x1, y1), (x1, y1 + line_length_y), color, thickness)  # top-left
+    cv2.line(img, (x1, y1), (x1 + line_length_x, y1), color, thickness)
 
-    # Filter by confidence
-    mask = scores > confidence_threshold
-    boxes = boxes[mask]
-    scores = scores[mask]
-    class_ids = class_ids[mask]
+    cv2.line(img, (x1, y2), (x1, y2 - line_length_y), color, thickness)  # bottom-left
+    cv2.line(img, (x1, y2), (x1 + line_length_x, y2), color, thickness)
 
-    # Convert boxes from [x, y, w, h] to [x1, y1, x2, y2]
-    boxes[:, 2:] += boxes[:, :2]
+    cv2.line(img, (x2, y1), (x2 - line_length_x, y1), color, thickness)  # top-right
+    cv2.line(img, (x2, y1), (x2, y1 + line_length_y), color, thickness)
 
-    # Scale boxes to original image size
-    h, w = original_shape
-    boxes[:, [0, 2]] *= w
-    boxes[:, [1, 3]] *= h
+    cv2.line(img, (x2, y2), (x2, y2 - line_length_y), color, thickness)  # bottom-right
+    cv2.line(img, (x2, y2), (x2 - line_length_x, y2), color, thickness)
 
-    # Apply NMS for each class separately
-    results = []
-    for class_id in np.unique(class_ids):
-        class_mask = class_ids == class_id
-        class_boxes = boxes[class_mask]
-        class_scores = scores[class_mask]
-        
-        indices = cv2.dnn.NMSBoxes(
-            class_boxes.tolist(), 
-            class_scores.tolist(), 
-            confidence_threshold, 
-            iou_threshold
-        )
-
-        for i in indices:
-            box = class_boxes[i]
-            score = class_scores[i]
-            x1, y1, x2, y2 = map(int, box)
-            results.append((x1, y1, x2, y2, float(score), int(class_id)))
-
-    return results
+    return img
 
-def process_image(image):
-    orig_image = image.copy()
-    processed_image, original_shape = preprocess_image(image)
-    
-    # Run inference
-    inputs = {ort_session.get_inputs()[0].name: processed_image}
-    outputs = ort_session.run(None, inputs)
-    
-    results = postprocess_results(outputs, original_shape)
+def process_frame(frame, license_plate_detector, vehicle_detector, ort_session):
+    # Detect vehicles
+    vehicle_results = vehicle_detector(frame, classes=[2, 3, 5, 7])  # cars, motorcycles, bus, trucks
     
-    # Draw bounding boxes on the image
-    for x1, y1, x2, y2, score, class_id in results:
-        # Draw rectangle with white color
-        cv2.rectangle(orig_image, (x1, y1), (x2, y2), (255, 255, 255), 2)
-        
-        # Get class name
-        class_name = CLASSES.get(class_id, f"Class_{class_id}")
-        label = f"{class_name}: {score:.2f}"
-        
-        # Add label background and text
-        (text_width, text_height), _ = cv2.getTextSize(
-            label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1
-        )
-        
-        # Draw black background for text
-        cv2.rectangle(
-            orig_image, 
-            (x1, y1 - text_height - 4), 
-            (x1 + text_width, y1), 
-            (0, 0, 0), 
-            -1
-        )
-        
-        # Draw white text
-        cv2.putText(
-            orig_image,
-            label,
-            (x1, y1 - 5),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.6,
-            (255, 255, 255),
-            1
-        )
+    # Process each vehicle
+    for vehicle in vehicle_results[0].boxes.data:
+        x1, y1, x2, y2, score, class_id = vehicle
+        if score > 0.5:  # Confidence threshold
+            # Draw vehicle border
+            draw_border(frame, 
+                       (int(x1), int(y1)), 
+                       (int(x2), int(y2)), 
+                       color=(0, 255, 0), 
+                       thickness=25,
+                       line_length_x=200, 
+                       line_length_y=200)
+            
+            # Detect license plate in vehicle region
+            vehicle_crop = frame[int(y1):int(y2), int(x1):int(x2)]
+            license_results = license_plate_detector(vehicle_crop)
+            
+            for license_plate in license_results[0].boxes.data:
+                lp_x1, lp_y1, lp_x2, lp_y2, lp_score, _ = license_plate
+                if lp_score > 0.5:
+                    # Adjust coordinates to full frame
+                    abs_lp_x1 = int(x1 + lp_x1)
+                    abs_lp_y1 = int(y1 + lp_y1)
+                    abs_lp_x2 = int(x1 + lp_x2)
+                    abs_lp_y2 = int(y1 + lp_y2)
+                    
+                    # Draw license plate box
+                    cv2.rectangle(frame, 
+                                (abs_lp_x1, abs_lp_y1), 
+                                (abs_lp_x2, abs_lp_y2), 
+                                (0, 0, 255), 12)
+                    
+                    # Extract and process license plate for OCR
+                    license_crop = frame[abs_lp_y1:abs_lp_y2, abs_lp_x1:abs_lp_x2]
+                    if license_crop.size > 0:
+                        # Prepare license crop for ONNX model
+                        license_crop_resized = cv2.resize(license_crop, (640, 640))
+                        license_crop_processed = np.transpose(license_crop_resized, (2, 0, 1)).astype(np.float32) / 255.0
+                        license_crop_processed = np.expand_dims(license_crop_processed, axis=0)
+                        
+                        # Run OCR inference
+                        try:
+                            inputs = {ort_session.get_inputs()[0].name: license_crop_processed}
+                            outputs = ort_session.run(None, inputs)
+                            
+                            # Process OCR output (adjust based on your model's output format)
+                            # This is a placeholder - adjust based on your ONNX model's output
+                            license_number = "ABC123"  # Replace with actual OCR processing
+                            
+                            # Display license plate number
+                            H, W, _ = license_crop.shape
+                            license_crop_display = cv2.resize(license_crop, (int(W * 400 / H), 400))
+                            
+                            try:
+                                # Display license crop and number above vehicle
+                                h_crop, w_crop, _ = license_crop_display.shape
+                                center_x = int((x1 + x2) / 2)
+                                
+                                # Display license plate crop
+                                frame[int(y1) - h_crop - 100:int(y1) - 100,
+                                      int(center_x - w_crop/2):int(center_x + w_crop/2)] = license_crop_display
+                                
+                                # White background for text
+                                cv2.rectangle(frame,
+                                            (int(center_x - w_crop/2), int(y1) - h_crop - 400),
+                                            (int(center_x + w_crop/2), int(y1) - h_crop - 100),
+                                            (255, 255, 255),
+                                            -1)
+                                
+                                # Draw license number
+                                (text_width, text_height), _ = cv2.getTextSize(
+                                    license_number,
+                                    cv2.FONT_HERSHEY_SIMPLEX,
+                                    4.3,
+                                    17)
+                                
+                                cv2.putText(frame,
+                                            license_number,
+                                            (int(center_x - text_width/2), int(y1 - h_crop - 250 + text_height/2)),
+                                            cv2.FONT_HERSHEY_SIMPLEX,
+                                            4.3,
+                                            (0, 0, 0),
+                                            17)
+                            except Exception as e:
+                                st.error(f"Error displaying results: {str(e)}")
+                        except Exception as e:
+                            st.error(f"Error in OCR processing: {str(e)}")
     
-    return cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
+    return frame
 
-def process_video(video_path):
+def process_video(video_path, license_plate_detector, vehicle_detector, ort_session):
     cap = cv2.VideoCapture(video_path)
     
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -138,12 +136,10 @@ def process_video(video_path):
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-    out = cv2.VideoWriter(
-        temp_file.name,
-        cv2.VideoWriter_fourcc(*'mp4v'),
-        fps,
-        (width, height)
-    )
+    out = cv2.VideoWriter(temp_file.name,
+                         cv2.VideoWriter_fourcc(*'mp4v'),
+                         fps,
+                         (width, height))
     
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     progress_bar = st.progress(0)
@@ -154,8 +150,8 @@ def process_video(video_path):
         if not ret:
             break
         
-        processed_frame = process_image(frame)
-        out.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
+        processed_frame = process_frame(frame, license_plate_detector, vehicle_detector, ort_session)
+        out.write(processed_frame)
         
         frame_count += 1
         progress_bar.progress(frame_count / total_frames)
@@ -167,43 +163,37 @@ def process_video(video_path):
     return temp_file.name
 
 # Streamlit UI
-st.title("Vehicle and License Plate Detection")
-
-# Add confidence threshold slider
-confidence_threshold = st.slider(
-    "Confidence Threshold",
-    min_value=0.0,
-    max_value=1.0,
-    value=0.25,
-    step=0.05
-)
+st.title("Advanced Vehicle and License Plate Detection")
 
-uploaded_file = st.file_uploader(
-    "Choose an image or video file",
-    type=["jpg", "jpeg", "png", "mp4"]
-)
-
-if uploaded_file is not None:
-    file_type = uploaded_file.type.split('/')[0]
-    
-    if file_type == "image":
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-        
-        if st.button("Detect Objects"):
-            with st.spinner("Processing image..."):
-                processed_image = process_image(np.array(image))
-                st.image(processed_image, caption="Processed Image", use_column_width=True)
+try:
+    license_plate_detector, vehicle_detector, ort_session = load_models()
     
-    elif file_type == "video":
-        tfile = tempfile.NamedTemporaryFile(delete=False)
-        tfile.write(uploaded_file.read())
+    uploaded_file = st.file_uploader("Choose an image or video file", type=["jpg", "jpeg", "png", "mp4"])
+
+    if uploaded_file is not None:
+        file_type = uploaded_file.type.split('/')[0]
         
-        st.video(tfile.name)
+        if file_type == "image":
+            image = Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+            
+            if st.button("Detect"):
+                with st.spinner("Processing image..."):
+                    # Convert PIL Image to CV2 format
+                    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+                    processed_image = process_frame(image_cv, license_plate_detector, vehicle_detector, ort_session)
+                    processed_image = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+                    st.image(processed_image, caption="Processed Image", use_column_width=True)
         
-        if st.button("Detect Objects"):
-            with st.spinner("Processing video..."):
-                processed_video = process_video(tfile.name)
-                st.video(processed_video)
-
-st.write("Upload an image or video to detect vehicles and license plates.")
+        elif file_type == "video":
+            tfile = tempfile.NamedTemporaryFile(delete=False)
+            tfile.write(uploaded_file.read())
+            st.video(tfile.name)
+            
+            if st.button("Detect"):
+                with st.spinner("Processing video..."):
+                    processed_video = process_video(tfile.name, license_plate_detector, vehicle_detector, ort_session)
+                    st.video(processed_video)
+
+except Exception as e:
+    st.error(f"Error loading models: {str(e)}")