Create processor.py

processor.py

@@ -0,0 +1,294 @@
+# processor.py
+
+import cv2
+import numpy as np
+import smtplib
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
+from email.mime.base import MIMEBase
+from email import encoders
+import os
+from ultralytics import YOLO
+from transformers import AutoModel, AutoProcessor
+from PIL import Image, ImageDraw, ImageFont
+import re
+import torch
+
+# Email credentials (replace with your App-Specific Password)
+FROM_EMAIL = "[email protected]"
+EMAIL_PASSWORD = "cawxqifzqiwjufde"  # Use App-Specific Password here
+TO_EMAIL = "[email protected]"
+SMTP_SERVER = 'smtp.gmail.com'
+SMTP_PORT = 465
+
+# Arabic dictionary for converting license plate text
+arabic_dict = {
+    "0": "٠", "1": "١", "2": "٢", "3": "٣", "4": "٤", "5": "٥",
+    "6": "٦", "7": "٧", "8": "٨", "9": "٩", "A": "ا", "B": "ب",
+    "J": "ح", "D": "د", "R": "ر", "S": "س", "X": "ص", "T": "ط",
+    "E": "ع", "G": "ق", "K": "ك", "L": "ل", "Z": "م", "N": "ن",
+    "H": "ه", "U": "و", "V": "ي", " ": " "
+}
+class_colors = {
+    0: (0, 255, 0),    # Green (Helmet)
+    1: (255, 0, 0),    # Blue (License Plate)
+    2: (0, 0, 255),    # Red (MotorbikeDelivery)
+    3: (255, 255, 0),  # Cyan (MotorbikeSport)
+    4: (255, 0, 255),  # Magenta (No Helmet)
+    5: (0, 255, 255),  # Yellow (Person)
+}
+
+# Load the OCR model
+processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True)
+model_ocr = AutoModel.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True).to('cuda')
+
+# Load YOLO model
+model = YOLO('yolov8_Medium.pt')  # Update the path as needed
+
+# Define lane area coordinates (example coordinates)
+red_lane = np.array([[2,1583],[1,1131],[1828,1141],[1912,1580]], np.int32)
+
+# Dictionary to track violations per license plate
+violations_dict = {}
+
+def filter_license_plate_text(license_plate_text):
+    license_plate_text = re.sub(r'[^A-Z0-9]+', "", license_plate_text)
+    match = re.search(r'(\d{4})\s*([A-Z]{2})', license_plate_text)
+    return f"{match.group(1)} {match.group(2)}" if match else None
+
+def convert_to_arabic(license_plate_text):
+    return "".join(arabic_dict.get(char, char) for char in license_plate_text)
+
+def send_email(license_text, violation_image_path, violation_type): 
+    # Define the subject and body based on violation type
+    if violation_type == 'No Helmet, In Red Lane':
+        subject = 'تنبيه مخالفة: عدم ارتداء خوذة ودخول المسار الأيسر'
+        body = f"لعدم ارتداء الخوذة ولدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
+        
+    elif violation_type == 'In Red Lane':
+        subject = 'تنبيه مخالفة: دخول المسار الأيسر'
+        body = f"لدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
+    else:
+        violation_type == 'No Helmet'
+        subject = 'تنبيه مخالفة: عدم ارتداء خوذة'
+        body = f"لعدم ارتداء الخوذة ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
+        
+    # Create the email message
+    msg = MIMEMultipart()
+    msg['From'] = FROM_EMAIL
+    msg['To'] = TO_EMAIL
+    msg['Subject'] = subject 
+    msg.attach(MIMEText(body, 'plain'))
+
+    # Attach the violation image
+    if os.path.exists(violation_image_path):
+        with open(violation_image_path, 'rb') as attachment_file:
+            part = MIMEBase('application', 'octet-stream')
+            part.set_payload(attachment_file.read())
+            encoders.encode_base64(part)
+            part.add_header('Content-Disposition', f'attachment; filename={os.path.basename(violation_image_path)}')
+            msg.attach(part)
+
+    # Send the email using SMTP
+    try:
+        with smtplib.SMTP_SSL(SMTP_SERVER, SMTP_PORT) as server:
+            server.login(FROM_EMAIL, EMAIL_PASSWORD)
+            server.sendmail(FROM_EMAIL, TO_EMAIL, msg.as_string())
+            print("Email with attachment sent successfully!")
+    except Exception as e:
+        print(f"Failed to send email: {e}")
+
+def draw_text_pil(img, text, position, font_path, font_size, color):
+    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    
+    draw = ImageDraw.Draw(img_pil)
+    
+    try:
+        font = ImageFont.truetype(font_path, size=font_size)
+    except IOError:
+        print(f"Font file not found at {font_path}. Using default font.")
+        font = ImageFont.load_default()
+        
+    draw.text(position, text, font=font, fill=color)
+    
+    img_np = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+    return img_np
+
+def process_video(video_path, font_path, violation_image_path='violation.jpg'):
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        print("Error opening video file")
+        return None
+    
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    output_video_path = '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))  
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+    
+    margin_y = 50
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break  
+
+        cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)  
+
+        results = model.track(frame)
+        
+        for box in results[0].boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())  
+            label = model.names[int(box.cls)]  
+            color = class_colors.get(int(box.cls), (255, 255, 255))
+            confidence = box.conf[0].item()
+
+            helmet_violation = False
+            lane_violation = False
+            violation_type = []
+
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)  
+            cv2.putText(frame, f'{label}: {confidence:.2f}', (x1, y1 - 10), 
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+
+            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")
+
+                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())  
+                    sub_label = model.names[int(result.cls)]
+                    sub_color = (255, 0, 0)  
+
+                    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]
+
+                        if helmet_violation or lane_violation:
+                            cv2.imwrite(violation_image_path, frame)
+                            license_plate_pil = Image.fromarray(cv2.cvtColor(license_crop, cv2.COLOR_BGR2RGB))
+                            temp_image_path = 'license_plate.png'
+                            license_plate_pil.save(temp_image_path)
+                            # Placeholder for OCR
+                            # license_plate_text = model_ocr.chat(processor, temp_image_path, ocr_type='ocr')
+                            # For demonstration, we'll mock the OCR result
+                            license_plate_text = "1234AB"
+                            filtered_text = filter_license_plate_text(license_plate_text)
+
+                            if filtered_text:
+                                if filtered_text not in violations_dict:
+                                    violations_dict[filtered_text] = violation_type  
+                                    send_email(filtered_text, violation_image_path, ', '.join(violation_type)) 
+                                else:
+                                    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)) 
+
+                                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))
+
+        out.write(frame)
+    
+    cap.release()
+    out.release()
+    return output_video_path
+
+def process_image(image_path, font_path, violation_image_path='violation.jpg'):
+    frame = cv2.imread(image_path)
+    if frame is None:
+        print("Error loading image")
+        return None
+
+    cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)  
+
+    results = model.track(frame)
+    
+    for box in results[0].boxes:
+        x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())  
+        label = model.names[int(box.cls)]  
+        color = class_colors.get(int(box.cls), (255, 255, 255))
+        confidence = box.conf[0].item()
+
+        helmet_violation = False
+        lane_violation = False
+        violation_type = []
+
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)  
+        cv2.putText(frame, f'{label}: {confidence:.2f}', (x1, y1 - 10), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+
+        if label == 'MotorbikeDelivery' and confidence >= 0.4:
+            motorbike_crop = frame[max(0, y1 - 50):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")
+
+            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())  
+                sub_label = model.names[int(result.cls)]
+                sub_color = (255, 0, 0)  
+
+                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]
+
+                    if helmet_violation or lane_violation:
+                        cv2.imwrite(violation_image_path, frame)
+                        license_plate_pil = Image.fromarray(cv2.cvtColor(license_crop, cv2.COLOR_BGR2RGB))
+                        temp_image_path = 'license_plate.png'
+                        license_plate_pil.save(temp_image_path)
+                        # Placeholder for OCR
+                        # license_plate_text = model_ocr.chat(processor, temp_image_path, ocr_type='ocr')
+                        # For demonstration, we'll mock the OCR result
+                        license_plate_text = "1234AB"
+                        filtered_text = filter_license_plate_text(license_plate_text)
+
+                        if filtered_text:
+                            if filtered_text not in violations_dict:
+                                violations_dict[filtered_text] = violation_type  
+                                send_email(filtered_text, violation_image_path, ', '.join(violation_type)) 
+                            else:
+                                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)) 
+
+                            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))
+
+    return frame