Update processor.py

processor.py

@@ -2,26 +2,25 @@
 
 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
+import smtplib
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
+from email.mime.base import MIMEBase
+from email import encoders
 
-# Email credentials (replace with your App-Specific Password)
-FROM_EMAIL = "Fares5675@gmail.com"
-EMAIL_PASSWORD = "cawxqifzqiwjufde"  # Use App-Specific Password here
-TO_EMAIL = "kduwaysan@gmail.com"
+# Email credentials (Use environment variables for security)
+FROM_EMAIL = os.getenv("FROM_EMAIL")
+EMAIL_PASSWORD = os.getenv("EMAIL_PASSWORD")
+TO_EMAIL = os.getenv("TO_EMAIL")
 SMTP_SERVER = 'smtp.gmail.com'
 SMTP_PORT = 465
 
-# Arabic dictionary for converting license plate text
+# Arabic dictionary
 arabic_dict = {
     "0": "٠", "1": "١", "2": "٢", "3": "٣", "4": "٤", "5": "٥",
     "6": "٦", "7": "٧", "8": "٨", "9": "٩", "A": "ا", "B": "ب",
@@ -38,49 +37,44 @@ class_colors = {
     5: (0, 255, 255),  # Yellow (Person)
 }
 
-# Load the OCR model
+# Load models
 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')
+model = YOLO('yolov8_Medium.pt')  # Update path as needed
 
-# Load YOLO model
-model = YOLO('yolov8_Medium.pt')  # Update the path as needed
-
-# Define lane area coordinates (example coordinates)
+# Define lane area
 red_lane = np.array([[2,1583],[1,1131],[1828,1141],[1912,1580]], np.int32)
 
-# Dictionary to track violations per license plate
+# Violation tracking
 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)
+def filter_license_plate_text(text):
+    text = re.sub(r'[^A-Z0-9]+', "", text)
+    match = re.search(r'(\d{4})\s*([A-Z]{2})', 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 convert_to_arabic(text):
+    return "".join(arabic_dict.get(char, char) for char in text)
+
+def send_email(license_text, violation_image_path, violation_type):
+    subject = {
+        'No Helmet, In Red Lane': 'تنبيه مخالفة: عدم ارتداء خوذة ودخول المسار الأيسر',
+        'In Red Lane': 'تنبيه مخالفة: دخول المسار الأيسر',
+        'No Helmet': 'تنبيه مخالفة: عدم ارتداء خوذة'
+    }.get(violation_type, 'تنبيه مخالفة')
+
+    body = {
+        'No Helmet, In Red Lane': f"لعدم ارتداء الخوذة ولدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة",
+        'In Red Lane': f"لدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة",
+        'No Helmet': f"لعدم ارتداء الخوذة ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
+    }.get(violation_type, f"تم تغريم دراجة نارية التي تحمل لوحة ({license_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['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')
@@ -89,200 +83,65 @@ def send_email(license_text, violation_image_path, violation_type):
             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)
-                            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:
-                                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)  
+    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
 
+def process_frame(frame, font_path, violation_image_path='violation.jpg'):
     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)]  
+        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)
+            violation_types = []
             if in_red_lane >= 0:
-                lane_violation = True
-                violation_type.append("In Red Lane")
+                violation_types.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)
-
+            for sub_box in sub_results[0].boxes:
+                sub_x1, sub_y1, sub_x2, sub_y2 = map(int, sub_box.xyxy[0].cpu().numpy())
+                sub_label = model.names[int(sub_box.cls)]
                 if sub_label == 'No_Helmet':
-                    helmet_violation = True
-                    violation_type.append("No Helmet")
-                    continue
-                if sub_label == 'License_plate':
+                    violation_types.append("No Helmet")
+                elif sub_label == 'License_plate':
                     license_crop = motorbike_crop[sub_y1:sub_y2, sub_x1:sub_x2]
-
-                    if helmet_violation or lane_violation:
+                    if violation_types:
                         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)
+                        license_plate_pil.save('license_plate.png')
                         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:
                             if filtered_text not in violations_dict:
-                                violations_dict[filtered_text] = violation_type  
-                                send_email(filtered_text, violation_image_path, ', '.join(violation_type)) 
+                                violations_dict[filtered_text] = violation_types
+                                send_email(filtered_text, violation_image_path, ', '.join(violation_types))
                             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)) 
-
+                                current = set(violations_dict[filtered_text])
+                                new = set(violation_types)
+                                updated = current | new
+                                if updated != current:
+                                    violations_dict[filtered_text] = list(updated)
+                                    send_email(filtered_text, violation_image_path, ', '.join(updated))
                             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
+                            frame = draw_text_pil(frame, filtered_text, (x1, y2 + 30), font_path, 30, (255, 255, 255))
+                            frame = draw_text_pil(frame, arabic_text, (x1, y2 + 60), font_path, 30, (0, 255, 0))
+    return frame