image_utils.py

import base64
import cv2
import numpy as np
from io import BytesIO
from PIL import Image

THUMBNAIL_SIZE = 160


def _b64_to_pil(base64_string):
    header, _, data = base64_string.partition(",")
    raw = base64.b64decode(data if data else base64_string)
    return Image.open(BytesIO(raw)).convert("RGB")


def _pil_to_cv2(image):
    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)


def _cv2_to_pil(bgr_image):
    return Image.fromarray(cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB))


def _crop_face_thumbnail(image, bbox, size):
    left, top, right, bottom = (int(c) for c in bbox)
    padding = max(int(max(right - left, bottom - top) * 0.35), 10)
    padded_left = max(0, left - padding)
    padded_top = max(0, top - padding)
    padded_right = min(image.width, right + padding)
    padded_bottom = min(image.height, bottom + padding)
    return image.crop((padded_left, padded_top, padded_right, padded_bottom)).resize(
        (size, size), Image.LANCZOS
    )


def _thumbnail_to_base64(thumbnail):
    buffer = BytesIO()
    thumbnail.save(buffer, format="JPEG", quality=88)
    return f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode()}"


def _sort_faces_left_to_right(faces):
    return sorted(faces, key=lambda face: float(face.bbox[0]))


def _face_to_dict(face_index, face, image):
    thumbnail = _crop_face_thumbnail(image, face.bbox, THUMBNAIL_SIZE)
    return {
        "idx": face_index,
        "thumbnail": _thumbnail_to_base64(thumbnail),
        "bbox": [int(c) for c in face.bbox],
        "det_score": round(float(face.det_score), 3),
    }


def _pil_to_b64(image):
    buffer = BytesIO()
    image.save(buffer, format="JPEG", quality=88)
    return f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode()}"


def _collect_source_images(selected_indices, sources):
    source_images = []
    for face_index in selected_indices:
        source = sources.get(str(face_index))
        if source:
            try:
                source_images.append(_b64_to_pil(source))
            except Exception:
                pass
    return source_images