utils.py

import base64
import io
from typing import Dict, List

import numpy as np
from PIL import Image


def pil_to_base64(img: Image.Image, max_size: int = 1920, quality: int = 90) -> str:
    """
    Convert PIL image to base64. 
    Increased max_size to 1920 to ensure Editor precision.
    """
    img = img.copy()
    img.thumbnail((max_size, max_size)) # Was 800, changed to 1920
    buffer = io.BytesIO()
    img.save(buffer, format="JPEG", quality=quality)
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode("utf-8")

def mask_to_coco_rle(mask: np.ndarray) -> Dict:
    """Convert binary mask to COCO RLE format."""
    from pycocotools import mask as coco_mask

    mask_uint8 = (mask.astype(bool) * 255).astype(np.uint8)
    fortran_mask = np.asfortranarray(mask_uint8)
    rle = coco_mask.encode(fortran_mask)
    # pycocotools returns bytes; COCO JSON needs UTF-8 string
    rle["counts"] = rle["counts"].decode("utf-8")
    return rle


def mask_to_polygon(mask: np.ndarray) -> List[List[float]]:
    """Convert binary mask to polygon coordinates (COCO-style)."""
    import cv2

    mask_uint8 = (mask.astype(bool) * 255).astype(np.uint8)
    contours, _ = cv2.findContours(
        mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    polygons: List[List[float]] = []
    for contour in contours:
        if len(contour) >= 3:
            polygon = contour.flatten().astype(float).tolist()
            # Need at least 3 points = 6 values
            if len(polygon) >= 6:
                polygons.append(polygon)

    return polygons if polygons else [[]]


def bbox_to_yolo(bbox: List[float], img_width: float, img_height: float) -> List[float]:
    """Convert [x1, y1, x2, y2] to YOLO [cx, cy, w, h] normalized to [0,1]."""
    x1, y1, x2, y2 = bbox
    center_x = ((x1 + x2) / 2.0) / img_width
    center_y = ((y1 + y2) / 2.0) / img_height
    width = (x2 - x1) / img_width
    height = (y2 - y1) / img_height
    return [center_x, center_y, width, height]


def mask_to_yolo_segmentation(mask: np.ndarray, img_width: float, img_height: float) -> List[float]:
    """Convert mask to YOLO segmentation format (normalized polygon)."""
    import cv2

    mask_uint8 = (mask.astype(bool) * 255).astype(np.uint8)
    contours, _ = cv2.findContours(
        mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    if not contours:
        return []

    # Use the largest contour by area
    contour = max(contours, key=cv2.contourArea)
    polygon = contour.flatten().astype(float).tolist()

    normalized = [
        coord / (img_width if idx % 2 == 0 else img_height)
        for idx, coord in enumerate(polygon)
    ]
    return normalized