Update app.py

app.py

@@ -46,6 +46,8 @@ class_colors = {
 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')
 
+# Define lane area coordinates (example coordinates)
+red_lane = np.array([[2, 1583], [1, 1131], [1828, 1141], [1912, 1580]], np.int32)
 
 # YOLO inference function
 def run_yolo(image):
@@ -90,46 +92,143 @@ def process_image(uploaded_file):
 
 # Process and save uploaded videos
 @st.cache_data
+# Define the function to process the video
 def process_video_and_save(uploaded_file):
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
-        temp_file.write(uploaded_file.read())
-        temp_file_path = temp_file.name
+    # Path for Arabic font
+    font_path = "/kaggle/input/fontss/alfont_com_arial-1.ttf"
 
-    video = cv2.VideoCapture(temp_file_path)
-    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-    frames = []
-    current_frame = 0
-    start_time = time.time()
+    # Paths for saving violation images
+    violation_image_path = '/kaggle/working/violation.jpg'
 
-    progress_bar = st.progress(0)
-    progress_text = st.empty()
+    # Track emails already sent to avoid duplicate emails
+    sent_emails = {}
 
-    while True:
-        ret, frame = video.read()
-        if not ret:
-            break
-        results = run_yolo(frame)
-        processed_frame = process_results(results, frame)
-        processed_frame_rgb = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)
-        frames.append(processed_frame_rgb)
+    # Dictionary to track violations per license plate
+    violations_dict = {}
 
-        current_frame += 1
-        progress_percentage = int((current_frame / total_frames) * 100)
-        progress_bar.progress(progress_percentage)
-        progress_text.text(f"Processing frame {current_frame}/{total_frames} ({progress_percentage}%)")
+    # Video path (input)
+    video_path = "/kaggle/working/uploaded_video.mp4"  # Save the uploaded video file to this path
+    with open(video_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
 
-    video.release()
-    output_path = 'processed_video.mp4'
-    height, width, _ = frames[0].shape
-    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
+    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 = '/kaggle/working/output_violation.mp4'
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    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))
 
-    for frame in frames:
-        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-        out.write(frame_bgr)
+    margin_y = 50
 
+    # Process the video frame by frame
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break  # End of video
+
+        # Draw the red lane rectangle 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 = (255, 0, 0)  # Use a fixed color for bounding boxes
+            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)  # 3 is the thickness of the rectangle
+
+            # 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)
+
+            # 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)
+                            # Check if the license plate is already detected and saved
+                            if filtered_text:
+                                # Get the email from the database
+                                email = get_vehicle_information(filtered_text, lane_violation, helmet_violation, violation_image_path, "Riyadh")
+                            
+                                # Add the license plate and its violations to the violations dictionary
+                                if filtered_text not in violations_dict:
+                                    violations_dict[filtered_text] = violation_type  #{"1234AB":[no_Helmet,In_red_Lane]}
+                                    send_email(filtered_text, violation_image_path, ', '.join(violation_type), email)
+                                else:
+                                    # Update the violations for the license plate if new ones are found
+                                    current_violations = set(violations_dict[filtered_text])  # no helmet
+                                    new_violations = set(violation_type)  # red lane, no helmet
+                                    updated_violations = list(current_violations | new_violations)  # red_lane, no helmet
+
+                                    # If new violations are found, update and send email
+                                    if updated_violations != violations_dict[filtered_text]:
+                                        violations_dict[filtered_text] = updated_violations
+                                        send_email(filtered_text, violation_image_path, ', '.join(updated_violations), email)
+
+                            # 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_path
+
+return output_video_path  # Return the path of the processed video
+
 
 
 
@@ -206,6 +305,193 @@ def send_email(license_text, violation_image_path, violation_type):
         server.sendmail(FROM_EMAIL, TO_EMAIL, msg.as_string())
         print("Email with attachment sent successfully!")
 
+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
+
+import sqlite3
+from datetime import datetime
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import Column, String, Integer, Boolean
+
+def get_vehicle_information(detected_license_plate, lane_violation, no_helmet, image_link, city):
+    # Get current date and time
+    current_datetime = datetime.now()
+    current_date = current_datetime.strftime('%Y-%m-%d')
+    current_time = current_datetime.strftime('%H:%M:%S')
+
+    # Connect to SQLite database (motorbike_detections.db)
+    try:
+        motorbike_conn = sqlite3.connect('motorbike_detections.db')
+        motorbike_cursor = motorbike_conn.cursor()
+
+        # Check if the detection already exists
+        check_query = '''
+            SELECT DetectionID FROM MotorbikeDetections
+            WHERE LicensePlate = ? AND Date = ? AND Time = ? AND ImageLink = ?
+        '''
+        motorbike_cursor.execute(check_query, (detected_license_plate, current_date, current_time, image_link))
+        existing_detection = motorbike_cursor.fetchone()
+
+        if existing_detection:
+            print(f"Detection for license plate {detected_license_plate} at {current_date} {current_time} already exists.")
+        else:
+            # Insert the new detection record
+            insert_query = '''
+                INSERT INTO MotorbikeDetections (Date, Time, City, LicensePlate, LaneViolation, NoHelmet, ImageLink)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+            '''
+            motorbike_cursor.execute(insert_query, (
+                current_date,
+                current_time,
+                city,
+                detected_license_plate,
+                int(lane_violation),  # Convert boolean to integer (1 or 0)
+                int(no_helmet),        # Convert boolean to integer (1 or 0)
+                image_link
+            ))
+
+            # Commit the transaction
+            motorbike_conn.commit()
+            print(f"Detection data for license plate {detected_license_plate} inserted successfully.")
+
+    except sqlite3.IntegrityError as e:
+        print(f"Integrity Error: {e}. This detection may already exist.")
+    except sqlite3.Error as e:
+        print(f"An error occurred while inserting detection data: {e}")
+        motorbike_conn.rollback()
+    finally:
+        # Close the motorbike detections database connection
+        motorbike_conn.close()
+
+    # Retrieve email from 'vehicle_information.db'
+    try:
+        # Create an engine and session for SQLAlchemy
+        engine = create_engine('sqlite:///vehicle_information.db')
+        Session = sessionmaker(bind=engine)
+        session = Session()
+
+        # Query the VehicleInformation table for the detected license plate
+        vehicle_info = session.query(VehicleInformation).filter_by(license_plate=detected_license_plate).first()
+
+        if vehicle_info:
+            print(f"Email found for license plate {detected_license_plate}: {vehicle_info.email}")
+            return vehicle_info.email
+        else:
+            print(f"No vehicle information found for license plate {detected_license_plate}.")
+            return None
+
+    except Exception as e:
+        print(f"An error occurred while retrieving vehicle information: {e}")
+        return None
+
+    finally:
+        # Close the SQLAlchemy session
+        session.close()
+
+import sqlite3
+from datetime import datetime
+
+def setup_motorbike_detections_db():
+    # Connect to SQLite database (or create it if it doesn't exist)
+    conn = sqlite3.connect('motorbike_detections.db')
+    cursor = conn.cursor()
+
+    # Drop the old table if it exists, useful for restructuring
+    cursor.execute('DROP TABLE IF EXISTS MotorbikeDetections')
+
+    # Create the new table with a unique constraint on LicensePlate, Date, Time, and ImageLink
+    cursor.execute('''
+    CREATE TABLE IF NOT EXISTS MotorbikeDetections (
+        DetectionID INTEGER PRIMARY KEY AUTOINCREMENT,
+        Date DATE NOT NULL,
+        Time TIME NOT NULL,
+        City VARCHAR(100),
+        LicensePlate VARCHAR(100),
+        LaneViolation BOOLEAN NOT NULL,
+        NoHelmet BOOLEAN NOT NULL,
+        ImageLink VARCHAR(255),
+        UNIQUE(LicensePlate, Date, Time, ImageLink)
+    )
+    ''')
+
+    # Commit changes
+    conn.commit()
+
+    # Close the connection
+    conn.close()
+
+
+from sqlalchemy import create_engine, Column, String, Integer
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+# Define the base for model creation
+Base = declarative_base()
+
+# Define the VehicleInformation table using SQLAlchemy ORM
+class VehicleInformation(Base):
+    __tablename__ = 'VehicleInformation'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    license_plate = Column(String(100), unique=True, nullable=False)
+    email = Column(String(255), nullable=False)
+    phone_number = Column(String(15), nullable=False)
+    driver_id = Column(String(50), nullable=False)  # Government ID or identity number
+
+def setup_database():
+    # Create an SQLite database (this could be any database like PostgreSQL, MySQL, etc.)
+    engine = create_engine('sqlite:///vehicle_information.db')
+
+    # Drop existing tables and recreate them (for ensuring clean database on rerun)
+    Base.metadata.drop_all(engine)
+    Base.metadata.create_all(engine)
+
+    # Create a session to interact with the database
+    Session = sessionmaker(bind=engine)
+    session = Session()
+
+    # Insert some dummy data into the VehicleInformation table
+    vehicle_data_1 = VehicleInformation(
+        license_plate='1234 AB',
+        email='[email protected]',
+        phone_number='0559947203',
+        driver_id='ID1110000000'
+    )
+
+    vehicle_data_2 = VehicleInformation(
+        license_plate='3321 AR',
+        email='[email protected]',
+        phone_number='0539003545',
+        driver_id='ID2220000000'
+    )
+
+    # Add records to the session
+    session.add(vehicle_data_1)
+    session.add(vehicle_data_2)
+
+    # Commit the records to the database
+    session.commit()
+
+    # Query the table to confirm data
+    vehicles = session.query(VehicleInformation).all()
+
+    # Close the session
+    session.close()
 
 # Streamlit app main function
 def main():