app.py

import gradio as gr
from ultralytics import YOLO
import cv2
import os
import pymysql
import boto3
from io import BytesIO
import io
from PIL import Image
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

# Initialize AWS S3 client
aws_access_key = "AKIAXECLNGBK5SXL2CER"
aws_secret_key = "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"
aws_region = "eu-west-3"

# Initialize the S3 client using environment variables
s3 = boto3.client(
    's3',
    aws_access_key_id=aws_access_key,
    aws_secret_access_key=aws_secret_key,
    region_name=aws_region
)

S3_BUCKET_NAME = 'savingbuckett5'
S3_FOLDER = 'Video-Processing/'

# Load YOLO model from the local directory (ensure the model is uploaded to your Hugging Face space)
model = YOLO("./YOLO_Model_v5.pt")

RDS_HOST = "database-2.cnqamusmkwon.eu-north-1.rds.amazonaws.com"
RDS_PORT = 3306 
DB_USER = "root"
DB_PASSWORD = "mkmk162345"
DB_NAME = "traffic"

def get_connection():
    return pymysql.connect(
        host=RDS_HOST,
        port=RDS_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        database=DB_NAME,
        cursorclass=pymysql.cursors.DictCursor
    )

def increment_road(id_value, increment_value, is_in=True):
    try:
        connection = get_connection()
        with connection.cursor() as cursor:
            select_sql = "SELECT id, road_in, road_out, road_current FROM traffic_counter_road WHERE id = %s"
            cursor.execute(select_sql, (id_value,))
            result = cursor.fetchone()
        
        if result:
            with connection.cursor() as cursor:
                if is_in:
                    new_road_in = result['road_in'] + increment_value
                    new_road_current = new_road_in - result['road_out']
                    update_sql = """
                    UPDATE traffic_counter_road
                    SET road_in = %s, road_current = %s
                    WHERE id = %s
                    """
                    cursor.execute(update_sql, (new_road_in, new_road_current, id_value))
                else:
                    new_road_out = result['road_out'] + increment_value
                    new_road_current = result['road_in'] - new_road_out
                    update_sql = """
                    UPDATE traffic_counter_road
                    SET road_out = %s, road_current = %s
                    WHERE id = %s
                    """
                    cursor.execute(update_sql, (new_road_out, new_road_current, id_value))
            
            connection.commit()
    except pymysql.MySQLError as e:
        print(f"Error: {e}")
    finally:
        if connection:
            connection.close()

def upload_frame_to_s3(frame, frame_number):
    # Convert the OpenCV frame (BGR) to a PIL image (RGB)
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image)

    # Save the PIL image to an in-memory file
    buffer = BytesIO()
    pil_image.save(buffer, format="JPEG")
    buffer.seek(0)

    # Define the S3 object key (file name)
    s3_key = f"{S3_FOLDER}frame_{frame_number}.jpg"

    # Upload the image to S3
    s3.upload_fileobj(buffer, S3_BUCKET_NAME, s3_key)

    print(f"Uploaded frame {frame_number} to S3 at {S3_BUCKET_NAME}/{s3_key}")

def process_video(video_path, count_type):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("Error opening video file.")
    
    box = (1650, 900, 2816, 1500)  # Define the area for license plates
    counter = 0
    License_plate = set()
    class_names = ['License Plate', 'Car', 'Motorcycle', 'Truck']

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model.track(frame, persist=True)
        for result in results:
            for boxes in result.boxes:
                bbox = boxes.xyxy[0].cpu().numpy()
                class_id = int(boxes.cls[0].cpu().numpy())  
                conf = boxes.conf[0].cpu().numpy()

                id = int(boxes.id[0].cpu().numpy()) if boxes.id is not None else -1

                x1, y1, x2, y2 = map(int, bbox)
                cropped_object = frame[y1:y2, x1:x2]
                # cv2.rectangle(frame, (x1, y1), (x2, y2), (208, 38, 7), 3)
                # label = f'ID: {id}, class: {class_names[class_id]} Conf: {conf:.2f}'
                # cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (208, 38, 7), 2, cv2.LINE_AA)

                # Check if the object is in the defined box for the license plate
                if x1 >= box[0] and y1 >= box[1] and x2 <= box[2] and y2 <= box[3]:
                    if id not in License_plate:
                        License_plate.add(id)
                        if count_type == "in":
                            print("It's counting IN")
                            increment_road(1, 1, is_in=True)  # Update the road traffic database (count in)
                        elif count_type == "out":
                            print("It's counting OUT")
                            increment_road(1, 1, is_in=False)  # Update the road traffic database (count out)

                        print("It's now uploading")
                        upload_frame_to_s3(cropped_object, counter)  # Save cropped license plate to S3
                        counter += 1


def insert_data(license_value):
    try:
        connection = get_connection()
        with connection.cursor() as cursor:
            
            insert_sql = """
            INSERT INTO license_plates (license_plate)
            VALUES (%s)
            """
            cursor.execute(insert_sql, (license_value))
        
        connection.commit()
        
    except pymysql.MySQLError as e:
        print(f"Error: {e}")
    
    finally:
        if connection:
            connection.close()
            
# Gradio function for counting vehicles in
def count_in(video):
    process_video(video, count_type="in")
    return "Processed vehicles counting 'in' successfully."

# Gradio function for counting vehicles out
def count_out(video):
    process_video(video, count_type="out")
    return "Processed vehicles counting 'out' successfully."

# tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
# model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)

def ocr(image):
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
    # model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id).to(device)
    # # if isinstance(image, np.ndarray):
    # #     image = Image.fromarray(image)
    # # Save the image to a temporary file in /tmp directory
    # temp_image_path = "/tmp/temp_image.jpg"
    # image.save(temp_image_path, format='JPEG')
    # # Perform OCR on the image
    # res = model.chat(tokenizer, image, ocr_type='ocr')
    
    # # Return the extracted text
    # return res
    try:
        # Convert image to PIL Image if it's a NumPy array
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
        model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id).to(device)

        # Ensure the /tmp directory exists
        temp_dir = "/tmp"
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)

        # Save the image to a temporary file in /tmp directory
        temp_image_path = os.path.join(temp_dir, "temp_image.jpg")
        image.save(temp_image_path, format='JPEG')

        # Perform OCR on the image using the file path
        res = model.chat(tokenizer=tokenizer, image=temp_image_path, ocr_type='ocr')  # Pass the file path here

        output_text = tokenizer.decode(res[0], skip_special_tokens=True)
        return output_text
        # Return the extracted text
        # return res['text']  # Adjust this based on the actual return structure
    except Exception as e:
        return str(e)

# Create Gradio interfaces for two endpoints: count_in and count_out
iface_in = gr.Interface(
    fn=count_in,
    inputs="video",
    outputs=None,
    api_name="count_in",  # This explicitly sets the api_name
    title="YOLO Video Object Detection (Count In)",
    description="Upload a video to count vehicles 'in' and save frames to S3."
)

iface_out = gr.Interface(
    fn=count_out,
    inputs="video",
    outputs=None,
    api_name="count_out",  # This explicitly sets the api_name
    title="YOLO Video Object Detection (Count Out)",
    description="Upload a video to count vehicles 'out' and save frames to S3."
)

iface_ocr = gr.Interface(
    fn=ocr,
    inputs="image",
    # inputs=gr.Image(type="pil"),
    outputs="text",
    api_name="ocr",  # This explicitly sets the api_name
    title="OCR Image Text Extraction",
)

# Create a tabbed interface for both endpoints
iface = gr.TabbedInterface([iface_in, iface_out, iface_ocr], ["Count In", "Count Out", "OCR"])

# Launch the Gradio app
iface.launch()