Update processor.py

processor.py

@@ -196,31 +196,127 @@ def process_image(image_path, font_path, violation_image_path='violation.jpg'):
     processed = process_frame(frame, font_path, violation_image_path)
     return processed
 
-def process_video(video_path, font_path, violation_image_path='violation.jpg'):
-    """Process a video file and save the processed video."""
+def process_video(video_path):
+    # Paths for saving violation images
+    violation_image_path = '/kaggle/working/violation.jpg'
+
+    # Track emails already sent to avoid duplicate emails
+    sent_emails = {}
+
+    # Dictionary to track violations per license plate
+    violations_dict = {}
+
+    # Open video file
     cap = cv2.VideoCapture(video_path)
+
+    # Check if the video file opened successfully
     if not cap.isOpened():
         print("Error opening video file")
         return None
 
+    # Define codec and output video settings
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_video_path = 'output_violation.mp4'
+    output_video_path = '/kaggle/working/output_violation.mp4'
     fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # Frame width
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # Frame height
     out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
+    margin_y = 50
+
+    # Process the video frame by frame
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
-            break  
+            break  # End of video
+
+        # Draw the red lane polygon on each frame
+        cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)  # Red lane
+
+        # Perform detection using YOLO on the current frame
+        results = model.track(frame)
+
+        # Process each detection in the results
+        for box in results[0].boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())  # Bounding box coordinates
+            label = model.names[int(box.cls)]  # Class name (MotorbikeDelivery, Helmet, etc.)
+            color = class_colors[int(box.cls)]
+            confidence = box.conf[0].item()
+
+            # Initialize flags and variables for the violations
+            helmet_violation = False
+            lane_violation = False
+            violation_type = []
+
+            # Draw bounding box around detected object
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)
 
-        # Optionally draw red lane
-        cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)
+            # Add label to the box (e.g., 'MotorbikeDelivery')
+            cv2.putText(frame, f'{label}: {confidence:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
-        processed_frame = process_frame(frame, font_path, violation_image_path)
-        out.write(processed_frame)
+            # Detect MotorbikeDelivery
+            if label == 'MotorbikeDelivery' and confidence >= 0.4:
+                motorbike_crop = frame[max(0, y1 - margin_y):y2, x1:x2]
+                delivery_center = ((x1 + x2) // 2, (y2))
+                in_red_lane = cv2.pointPolygonTest(red_lane, delivery_center, False)
+                if in_red_lane >= 0:
+                    lane_violation = True
+                    violation_type.append("In Red Lane")
 
+                # Perform detection within the cropped motorbike region
+                sub_results = model(motorbike_crop)
+
+                for result in sub_results[0].boxes:
+                    sub_x1, sub_y1, sub_x2, sub_y2 = map(int, result.xyxy[0].cpu().numpy())  # Bounding box coordinates
+                    sub_label = model.names[int(result.cls)]
+                    sub_color = (255, 0, 0)  # Red color for the bounding box of sub-objects
+
+                    # Draw bounding box around sub-detected objects (No_Helmet, License_plate, etc.)
+                    cv2.rectangle(motorbike_crop, (sub_x1, sub_y1), (sub_x2, sub_y2), sub_color, 2)
+                    cv2.putText(motorbike_crop, sub_label, (sub_x1, sub_y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, sub_color, 2)
+
+                    if sub_label == 'No_Helmet':
+                        helmet_violation = True
+                        violation_type.append("No Helmet")
+                        continue
+                    if sub_label == 'License_plate':
+                        license_crop = motorbike_crop[sub_y1:sub_y2, sub_x1:sub_x2]
+
+                        # Apply OCR if a violation is detected
+                        if helmet_violation or lane_violation:
+                            # Perform OCR on the license plate
+                            cv2.imwrite(violation_image_path, frame)
+                            license_plate_pil = Image.fromarray(cv2.cvtColor(license_crop, cv2.COLOR_BGR2RGB))
+                            temp_image_path = '/kaggle/working/license_plate.png'
+                            license_plate_pil.save(temp_image_path)
+                            license_plate_text = model_ocr.chat(processor, temp_image_path, ocr_type='ocr')
+                            filtered_text = filter_license_plate_text(license_plate_text)
+                            
+                            if filtered_text:
+                                # Track violations for the license plate
+                                if filtered_text not in violations_dict:
+                                    violations_dict[filtered_text] = violation_type
+                                    send_email(filtered_text, violation_image_path, ', '.join(violation_type))
+                                else:
+                                    # Update violations if new ones are found
+                                    current_violations = set(violations_dict[filtered_text])
+                                    new_violations = set(violation_type)
+                                    updated_violations = list(current_violations | new_violations)
+
+                                    if updated_violations != violations_dict[filtered_text]:
+                                        violations_dict[filtered_text] = updated_violations
+                                        send_email(filtered_text, violation_image_path, ', '.join(updated_violations))
+
+                                # Draw OCR text (English and Arabic) on the original frame
+                                arabic_text = convert_to_arabic(filtered_text)
+                                frame = draw_text_pil(frame, filtered_text, (x1, y2 + 30), font_path, font_size=30, color=(255, 255, 255))
+                                frame = draw_text_pil(frame, arabic_text, (x1, y2 + 60), font_path, font_size=30, color=(0, 255, 0))
+
+        # Write the processed frame to the output video
+        out.write(frame)
+
+    # Release resources when done
     cap.release()
     out.release()
+
     return output_video_path