utils.py

"""
utils.py
--------
Image I/O and visualization helpers:
  * safe image loading / validation
  * colored prediction mask (dark gray / bright yellow)
  * semi-transparent magenta crop overlay
  * combined three-panel image + PNG export for download
"""

from __future__ import annotations

import io

import numpy as np
from PIL import Image, ImageDraw, ImageFont, UnidentifiedImageError

# --------------------------------------------------------------------------- #
# Visualization palette
# --------------------------------------------------------------------------- #
COLOR_BACKGROUND = (50, 50, 50)     # dark gray     -> Class 0 (No Crop)
COLOR_CROP = (255, 255, 0)          # bright yellow -> Class 1 (Crop)
OVERLAY_COLOR = (255, 0, 255)       # magenta overlay for crop regions
OVERLAY_ALPHA = 0.5                 # overlay transparency (0..1)


class ImageLoadError(Exception):
    """Raised when an uploaded image is invalid or corrupted."""


def load_image(file) -> Image.Image:
    """
    Safely load an uploaded file into an RGB PIL image.

    Raises ImageLoadError on invalid / corrupted input.
    """
    try:
        image = Image.open(file)
        image.load()  # force a full decode so corrupted files fail here
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise ImageLoadError(
            "The uploaded file is not a valid image or is corrupted."
        ) from exc

    return image.convert("RGB")


def create_color_mask(mask: np.ndarray) -> Image.Image:
    """Convert a {0, 1} mask into an RGB image (dark gray / bright yellow)."""
    h, w = mask.shape
    color = np.zeros((h, w, 3), dtype=np.uint8)
    color[mask == 0] = COLOR_BACKGROUND
    color[mask == 1] = COLOR_CROP
    return Image.fromarray(color, mode="RGB")


def create_overlay(original: Image.Image, mask: np.ndarray) -> Image.Image:
    """Blend a semi-transparent magenta layer over predicted crop regions."""
    base = np.asarray(original.convert("RGB"), dtype=np.float32)
    crop_pixels = mask == 1

    overlay = base.copy()
    overlay[crop_pixels] = (
        (1.0 - OVERLAY_ALPHA) * base[crop_pixels]
        + OVERLAY_ALPHA * np.array(OVERLAY_COLOR, dtype=np.float32)
    )
    return Image.fromarray(overlay.clip(0, 255).astype(np.uint8), mode="RGB")


def _load_font(size: int):
    """Load a default font, gracefully handling older Pillow versions."""
    try:
        return ImageFont.load_default(size=size)
    except TypeError:  # Pillow without the size argument
        return ImageFont.load_default()


def combine_panels(
    original: Image.Image,
    color_mask: Image.Image,
    overlay: Image.Image,
    gap: int = 12,
    label_height: int = 36,
) -> Image.Image:
    """
    Combine the three panels horizontally into a single labeled image.

    All panels are resized to identical dimensions (the original image size)
    before being concatenated.
    """
    target_size = original.size  # (w, h)
    panels = [
        ("Original", original.resize(target_size)),
        ("Predicted Mask", color_mask.resize(target_size)),
        ("Crop Overlay", overlay.resize(target_size)),
    ]

    w, h = target_size
    total_w = w * 3 + gap * 2
    total_h = h + label_height

    canvas = Image.new("RGB", (total_w, total_h), (255, 255, 255))
    draw = ImageDraw.Draw(canvas)
    font = _load_font(20)

    x = 0
    for label, img in panels:
        canvas.paste(img, (x, label_height))
        try:
            text_w = draw.textlength(label, font=font)
        except AttributeError:  # very old Pillow
            text_w = len(label) * 6
        draw.text(
            (x + (w - text_w) / 2, max(0, (label_height - 20) / 2)),
            label,
            fill=(0, 0, 0),
            font=font,
        )
        x += w + gap

    return canvas


def image_to_png_bytes(image: Image.Image) -> bytes:
    """Encode a PIL image as PNG bytes (used by the download button)."""
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()