crop_utils.py

import base64
import os
from io import BytesIO

import cv2
import gradio as gr
import numpy as np
import pyrebase
import requests
from openai import OpenAI
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO

from prompts import remove_unwanted_prompt


def get_middle_thumbnail(input_image: Image, grid_size=(10, 10), padding=3):
    """
    Extract the middle thumbnail from a sprite sheet, handling different aspect ratios
    and removing padding.

    Args:
        input_image: PIL Image
        grid_size: Tuple of (columns, rows)
        padding: Number of padding pixels on each side (default 3)

    Returns:
        PIL.Image: The middle thumbnail image with padding removed
    """
    sprite_sheet = input_image

    # Calculate thumbnail dimensions based on actual sprite sheet size
    sprite_width, sprite_height = sprite_sheet.size
    thumb_width_with_padding = sprite_width // grid_size[0]
    thumb_height_with_padding = sprite_height // grid_size[1]

    # Remove padding to get actual image dimensions
    thumb_width = thumb_width_with_padding - (2 * padding)  # 726 - 6 = 720
    thumb_height = thumb_height_with_padding - (2 * padding)  # varies based on input

    # Calculate the middle position
    total_thumbs = grid_size[0] * grid_size[1]
    middle_index = total_thumbs // 2

    # Calculate row and column of middle thumbnail
    middle_row = middle_index // grid_size[0]
    middle_col = middle_index % grid_size[0]

    # Calculate pixel coordinates for cropping, including padding offset
    left = (middle_col * thumb_width_with_padding) + padding
    top = (middle_row * thumb_height_with_padding) + padding
    right = left + thumb_width  # Don't add padding here
    bottom = top + thumb_height  # Don't add padding here

    # Crop and return the middle thumbnail
    middle_thumb = sprite_sheet.crop((left, top, right, bottom))
    return middle_thumb


def get_person_bbox(frame, model):
    """Detect person and return the largest bounding box"""
    results = model(frame, classes=[0])  # class 0 is person in COCO

    if not results or len(results[0].boxes) == 0:
        return None

    # Get all person boxes
    boxes = results[0].boxes.xyxy.cpu().numpy()
    # Calculate areas to find the largest person
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    largest_idx = np.argmax(areas)

    return boxes[largest_idx]


def generate_crops(frame):
    """Generate both 16:9 and 9:16 crops based on person detection"""
    # Load YOLO model
    model = YOLO("yolo11n.pt")

    # Convert PIL Image to cv2 format if needed
    if isinstance(frame, Image.Image):
        frame = cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR)

    original_height, original_width = frame.shape[:2]
    bbox = get_person_bbox(frame, model)

    if bbox is None:
        return None, None

    # Extract coordinates
    x1, y1, x2, y2 = map(int, bbox)
    person_height = y2 - y1
    person_width = x2 - x1
    person_center_x = (x1 + x2) // 2
    person_center_y = (y1 + y2) // 2

    # Generate 16:9 crop (focus on upper body)
    aspect_ratio_16_9 = 16 / 9
    crop_width_16_9 = min(original_width, int(person_height * aspect_ratio_16_9))
    crop_height_16_9 = min(original_height, int(crop_width_16_9 / aspect_ratio_16_9))

    # For 16:9, center horizontally and align top with person's top
    x1_16_9 = max(0, person_center_x - crop_width_16_9 // 2)
    x2_16_9 = min(original_width, x1_16_9 + crop_width_16_9)
    y1_16_9 = max(0, y1)  # Start from person's top
    y2_16_9 = min(original_height, y1_16_9 + crop_height_16_9)

    # Adjust if exceeding boundaries
    if x2_16_9 > original_width:
        x1_16_9 = original_width - crop_width_16_9
        x2_16_9 = original_width
    if y2_16_9 > original_height:
        y1_16_9 = original_height - crop_height_16_9
        y2_16_9 = original_height

    # Generate 9:16 crop (full body)
    aspect_ratio_9_16 = 9 / 16
    crop_width_9_16 = min(original_width, int(person_height * aspect_ratio_9_16))
    crop_height_9_16 = min(original_height, int(crop_width_9_16 / aspect_ratio_9_16))

    # For 9:16, center both horizontally and vertically
    x1_9_16 = max(0, person_center_x - crop_width_9_16 // 2)
    x2_9_16 = min(original_width, x1_9_16 + crop_width_9_16)
    y1_9_16 = max(0, person_center_y - crop_height_9_16 // 2)
    y2_9_16 = min(original_height, y1_9_16 + crop_height_9_16)

    # Adjust if exceeding boundaries
    if x2_9_16 > original_width:
        x1_9_16 = original_width - crop_width_9_16
        x2_9_16 = original_width
    if y2_9_16 > original_height:
        y1_9_16 = original_height - crop_height_9_16
        y2_9_16 = original_height

    # Create crops
    crop_16_9 = frame[y1_16_9:y2_16_9, x1_16_9:x2_16_9]
    crop_9_16 = frame[y1_9_16:y2_9_16, x1_9_16:x2_9_16]

    # Resize to standard dimensions
    crop_16_9 = cv2.resize(crop_16_9, (426, 240))  # 16:9 aspect ratio
    crop_9_16 = cv2.resize(crop_9_16, (240, 426))  # 9:16 aspect ratio

    return crop_16_9, crop_9_16


def visualize_crops(image, bbox, crops_info):
    """
    Visualize original bbox and calculated crops
    bbox: [x1, y1, x2, y2]
    crops_info: dict with 'crop_16_9' and 'crop_9_16' coordinates
    """
    viz = image.copy()

    # Draw original person bbox in blue
    cv2.rectangle(
        viz, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2
    )

    # Draw 16:9 crop in green
    crop_16_9 = crops_info["crop_16_9"]
    cv2.rectangle(
        viz,
        (int(crop_16_9["x1"]), int(crop_16_9["y1"])),
        (int(crop_16_9["x2"]), int(crop_16_9["y2"])),
        (0, 255, 0),
        2,
    )

    # Draw 9:16 crop in red
    crop_9_16 = crops_info["crop_9_16"]
    cv2.rectangle(
        viz,
        (int(crop_9_16["x1"]), int(crop_9_16["y1"])),
        (int(crop_9_16["x2"]), int(crop_9_16["y2"])),
        (0, 0, 255),
        2,
    )

    return viz


def encode_image_to_base64(image: Image.Image, format: str = "JPEG") -> str:
    """
    Convert a PIL image to a base64 string.

    Args:
        image: PIL Image object
        format: Image format to use for encoding (default: PNG)

    Returns:
        Base64 encoded string of the image
    """
    buffered = BytesIO()
    image.save(buffered, format=format)
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


def add_top_numbers(
    input_image,
    num_divisions=20,
    margin=90,
    font_size=120,
    dot_spacing=20,
):
    """
    Add numbered divisions across the top and bottom of any image with dotted vertical lines.

    Args:
        input_image (Image): PIL Image
        num_divisions (int): Number of divisions to create
        margin (int): Size of margin in pixels for numbers
        font_size (int): Font size for numbers
        dot_spacing (int): Spacing between dots in pixels
    """
    # Load the image
    original_image = input_image

    # Create new image with extra space for numbers on top and bottom
    new_width = original_image.width
    new_height = original_image.height + (
        2 * margin
    )  # Add margin to both top and bottom
    new_image = Image.new("RGB", (new_width, new_height), "white")

    # Paste original image in the middle
    new_image.paste(original_image, (0, margin))

    # Initialize drawing context
    draw = ImageDraw.Draw(new_image)

    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except OSError:
        print("Using default font")
        font = ImageFont.load_default(size=font_size)

    # Calculate division width
    division_width = original_image.width / num_divisions

    # Draw division numbers and dotted lines
    for i in range(num_divisions):
        x = (i * division_width) + (division_width / 2)

        # Draw number at top
        draw.text((x, margin // 2), str(i + 1), fill="black", font=font, anchor="mm")

        # Draw number at bottom
        draw.text(
            (x, new_height - (margin // 2)),
            str(i + 1),
            fill="black",
            font=font,
            anchor="mm",
        )

        # Draw dotted line from top margin to bottom margin
        y_start = margin
        y_end = new_height - margin

        # Draw dots with specified spacing
        current_y = y_start
        while current_y < y_end:
            draw.circle(
                [x - 1, current_y - 1, x + 1, current_y + 1],
                fill="black",
                width=5,
                radius=3,
            )
            current_y += dot_spacing

    return new_image


def crop_and_draw_divisions(
    input_image,
    left_division,
    right_division,
    num_divisions=20,
    line_color=(255, 0, 0),
    line_width=2,
    head_margin_percent=0.1,
):
    """
    Create both 9:16 and 16:9 crops and draw guide lines.

    Args:
        input_image (Image): PIL Image
        left_division (int): Left-side division number (1-20)
        right_division (int): Right-side division number (1-20)
        num_divisions (int): Total number of divisions (default=20)
        line_color (tuple): RGB color tuple for lines (default: red)
        line_width (int): Width of lines in pixels (default: 2)
        head_margin_percent (float): Percentage margin above head (default: 0.1)

    Returns:
        tuple: (cropped_image_16_9, image_with_lines, cropped_image_9_16)
    """
    yolo_model = YOLO("yolo11n.pt")
    # Calculate division width and boundaries
    division_width = input_image.width / num_divisions
    left_boundary = (left_division - 1) * division_width
    right_boundary = right_division * division_width

    # First get the 9:16 crop
    cropped_image_9_16 = input_image.crop(
        (left_boundary, 0, right_boundary, input_image.height)
    )

    # Run YOLO on the 9:16 crop to get person bbox
    bbox = yolo_model(cropped_image_9_16, classes=[0])[0].boxes.xyxy.cpu().numpy()[0]
    x1, y1, x2, y2 = bbox

    # Calculate top boundary with head margin
    head_margin = (y2 - y1) * head_margin_percent
    top_boundary = max(0, y1 - head_margin)

    # Calculate 16:9 dimensions based on the width between divisions
    crop_width = right_boundary - left_boundary
    crop_height_16_9 = int(crop_width * 9 / 16)

    # Calculate bottom boundary for 16:9
    bottom_boundary = min(input_image.height, top_boundary + crop_height_16_9)

    # Create 16:9 crop from original image
    cropped_image_16_9 = input_image.crop(
        (left_boundary, top_boundary, right_boundary, bottom_boundary)
    )

    # Draw guide lines for both crops on original image
    image_with_lines = input_image.copy()
    draw = ImageDraw.Draw(image_with_lines)

    # Draw vertical lines (for both crops)
    draw.line(
        [(left_boundary, 0), (left_boundary, input_image.height)],
        fill=line_color,
        width=line_width,
    )
    draw.line(
        [(right_boundary, 0), (right_boundary, input_image.height)],
        fill=line_color,
        width=line_width,
    )

    # Draw horizontal lines (for 16:9 crop)
    draw.line(
        [(left_boundary, top_boundary), (right_boundary, top_boundary)],
        fill=line_color,
        width=line_width,
    )
    draw.line(
        [(left_boundary, bottom_boundary), (right_boundary, bottom_boundary)],
        fill=line_color,
        width=line_width,
    )

    return cropped_image_16_9, image_with_lines, cropped_image_9_16


def analyze_image(numbered_input_image: Image, prompt, input_image):
    """
    Perform inference on an image using GPT-4V.

    Args:
        numbered_input_image (Image): PIL Image
        prompt (str): The prompt/question about the image
        input_image (Image): input image without numbers

    Returns:
        str: The model's response
    """
    client = OpenAI()
    base64_image = encode_image_to_base64(numbered_input_image, format="JPEG")

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                },
            ],
        }
    ]

    response = client.chat.completions.create(
        model="gpt-4o", messages=messages, max_tokens=300
    )

    messages.extend(
        [
            {"role": "assistant", "content": response.choices[0].message.content},
            {
                "role": "user",
                "content": "please return the response in the json with keys left_row and right_row",
            },
        ],
    )

    response = (
        client.chat.completions.create(model="gpt-4o", messages=messages)
        .choices[0]
        .message.content
    )

    left_index = response.find("{")
    right_index = response.rfind("}")

    try:
        if left_index != -1 and right_index != -1:
            response_json = eval(response[left_index : right_index + 1])
        cropped_image_16_9, image_with_lines, cropped_image_9_16 = (
            crop_and_draw_divisions(
                input_image=input_image,
                left_division=response_json["left_row"],
                right_division=response_json["right_row"],
            )
        )
    except Exception as e:
        print(e)
        return input_image, input_image, input_image

    return cropped_image_16_9, image_with_lines, cropped_image_9_16


def get_sprite_firebase(cid, rsid, uid):
    config = {
        "apiKey": f"{os.getenv('FIREBASE_API_KEY')}",
        "authDomain": f"{os.getenv('FIREBASE_AUTH_DOMAIN')}",
        "databaseURL": f"{os.getenv('FIREBASE_DATABASE_URL')}",
        "projectId": f"{os.getenv('FIREBASE_PROJECT_ID')}",
        "storageBucket": f"{os.getenv('FIREBASE_STORAGE_BUCKET')}",
        "messagingSenderId": f"{os.getenv('FIREBASE_MESSAGING_SENDER_ID')}",
        "appId": f"{os.getenv('FIREBASE_APP_ID')}",
        "measurementId": f"{os.getenv('FIREBASE_MEASUREMENT_ID')}",
    }

    firebase = pyrebase.initialize_app(config)
    db = firebase.database()
    account_id = os.getenv("ROLL_ACCOUNT")

    COLLAB_EDIT_LINK = "collab_sprite_link_handler"

    path = f"{account_id}/{COLLAB_EDIT_LINK}/{uid}/{cid}/{rsid}"

    data = db.child(path).get()
    return data.val()


def get_image_crop(cid=None, rsid=None, uid=None):
    """Function that returns both 16:9 and 9:16 crops"""
    image_paths = get_sprite_firebase(cid, rsid, uid)

    input_images = []
    mid_images = []
    cropped_image_16_9s = []
    images_with_lines = []
    cropped_image_9_16s = []

    for image_path in image_paths:
        response = requests.get(image_path)

        input_image = Image.open(BytesIO(response.content))
        input_images.append(input_image)

        # Get the middle thumbnail
        mid_image = get_middle_thumbnail(input_image)
        mid_images.append(mid_image)

        numbered_mid_image = add_top_numbers(
            input_image=mid_image,
            num_divisions=20,
            margin=50,
            font_size=30,
            dot_spacing=20,
        )

        cropped_image_16_9, image_with_lines, cropped_image_9_16 = analyze_image(
            numbered_mid_image, remove_unwanted_prompt(2), mid_image
        )
        cropped_image_16_9s.append(cropped_image_16_9)
        images_with_lines.append(image_with_lines)
        cropped_image_9_16s.append(cropped_image_9_16)

    return gr.Gallery(
        [
            *input_images,
            *mid_images,
            *cropped_image_16_9s,
            *images_with_lines,
            *cropped_image_9_16s,
        ]
    )