backend/predict.py

import base64
import hashlib
import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from dataclasses import dataclass
from typing import Tuple, List, Dict, Optional, Any, Callable

from transformers import AutoModelForImageSegmentation

from model import PokedexNet
from config import settings


@dataclass
class PredictionResult:
    pokemon_id: int
    name: str
    confidence: float
    detected_source: str
    top_5: List[Dict[str, Any]]
    debug_silhouette_b64: Optional[str] = None
    status: str = "CONFIDENT_POKEMON"
    votes: List[int] = None
    image_hash: str = ""
    is_ood_screen: bool = False
    ensemble_metrics: Dict[str, Any] = None
    original_width: int = 0
    original_height: int = 0


IS_SERVER = True

if IS_SERVER:
    DEVICE = "cpu"
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
else:
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ---------------------------------------------------------------------------
# Shared low-level helpers
# ---------------------------------------------------------------------------

def _largest_clean_component(
    mask: np.ndarray,
    min_area_ratio: float = 0.0,
) -> Optional[np.ndarray]:
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    if num_labels <= 1:
        return None
    areas = stats[1:, cv2.CC_STAT_AREA]
    largest_idx = int(np.argmax(areas)) + 1
    if areas[largest_idx - 1] < min_area_ratio * mask.size:
        return None
    return ((labels == largest_idx).astype(np.uint8)) * 255


# ---------------------------------------------------------------------------
# Mask quality & post-processing helpers
# ---------------------------------------------------------------------------

def _clean_mask(raw_mask: np.ndarray) -> np.ndarray:
    _, binary = cv2.threshold(raw_mask, 127, 255, cv2.THRESH_BINARY)

    k_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    closed  = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, k_close)

    clean = _largest_clean_component(closed)
    if clean is None:
        clean = closed

    _, final = cv2.threshold(clean, 127, 255, cv2.THRESH_BINARY)
    return final


def _center_and_pad_binary(mask: np.ndarray) -> np.ndarray:
    coords = cv2.findNonZero(mask)
    if coords is None:
        return np.full((128, 128), 255, dtype=np.uint8)

    x, y, w, h = cv2.boundingRect(coords)
    crop = mask[y:y + h, x:x + w]

    maior_lado    = max(w, h)
    margem        = int(maior_lado * 0.20)
    tamanho_final = maior_lado + margem * 2

    canvas = np.full((tamanho_final, tamanho_final), 255, dtype=np.uint8)
    y_off = (tamanho_final - h) // 2
    x_off = (tamanho_final - w) // 2
    canvas[y_off:y_off + h, x_off:x_off + w] = 255 - crop
    return canvas


def _audit_mask_quality(mask: np.ndarray, confidence_score: float) -> Tuple[bool, str]:
    if confidence_score < 0.30:
        return False, "BAD_CROP_LOW_CONFIDENCE"

    ratio = float(np.count_nonzero(mask)) / max(mask.size, 1)
    if ratio < 0.01:
        return False, "BAD_CROP_TOO_SMALL"
    if ratio > 0.90:
        return False, "BAD_CROP_TOO_LARGE"

    num_labels, _, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    if num_labels <= 1:
        return False, "BAD_CROP_TOO_SMALL"

    areas = stats[1:, cv2.CC_STAT_AREA]
    min_area_to_care = mask.size * 0.005
    significant_pieces = np.sum(areas > min_area_to_care)

    if significant_pieces > 8:
        return False, "BAD_CROP_TOO_FRAGMENTED"

    coords = cv2.findNonZero(mask)
    if coords is None:
        return False, "BAD_CROP_TOO_SMALL"

    bx, by, bw, bh = cv2.boundingRect(coords)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        return False, "BAD_CROP_TOO_SMALL"

    cnt = max(contours, key=cv2.contourArea)
    area = cv2.contourArea(cnt)
    hull = cv2.convexHull(cnt)
    hull_area = cv2.contourArea(hull)
    solidity = area / hull_area if hull_area > 0 else 0.0

    dist_map = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
    max_radius = np.max(dist_map)
    thickness_ratio = max_radius / max(bw, bh, 1)

    if ratio >= 0.03:
        min_solidity = 0.15
        min_thickness = 0.10
    else:
        min_solidity = 0.20
        min_thickness = 0.12

    if solidity < min_solidity or thickness_ratio < min_thickness:
        return False, "BAD_CROP_TOO_HOLLOW_OR_THIN"

    return True, "OK"


# ---------------------------------------------------------------------------
# BiRefNet model
# ---------------------------------------------------------------------------

birefnet_model: Any = None


def load_birefnet() -> Any:
    """Load BiRefNet model eagerly at startup."""
    global birefnet_model
    if birefnet_model is not None:
        return birefnet_model

    local_path = "/app/models/birefnet"
    if os.path.exists(local_path):
        try:
            birefnet_model = AutoModelForImageSegmentation.from_pretrained(
                local_path,
                trust_remote_code=True,
                local_files_only=True
            )
        except Exception:
            birefnet_model = None

    if birefnet_model is None:
        try:
            birefnet_model = AutoModelForImageSegmentation.from_pretrained(
                "ZhengPeng7/BiRefNet_lite",
                trust_remote_code=True,
                cache_dir="/app/models",
                local_files_only=True
            )
        except Exception as e:
            raise RuntimeError(
                f"Failed to load BiRefNet from local files. "
                f"Please verify that the model is downloaded in backend/models/birefnet. Error: {e}"
            )

    birefnet_model.to(DEVICE)
    birefnet_model.float()
    birefnet_model.eval()
    return birefnet_model


def _run_birefnet(img_rgb: np.ndarray, resolution: int) -> np.ndarray:
    pil = Image.fromarray(img_rgb)
    transform = transforms.Compose([
        transforms.Resize((resolution, resolution)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225]),
    ])
    tensor = transform(pil).unsqueeze(0).to(DEVICE).float()
    with torch.no_grad():
        pred = birefnet_model(tensor)[-1].sigmoid()
    mask_pred = pred[0, 0].cpu().numpy()
    return (mask_pred > 0.5).astype(np.uint8) * 255


def _extract_birefnet_silhouette_for_res(img_rgb: np.ndarray, resolution: int, w: int, h: int) -> Tuple[bool, np.ndarray, str, np.ndarray]:
    try:
        raw_mask = _run_birefnet(img_rgb, resolution)
    except Exception:
        raw_mask = np.zeros((resolution, resolution), dtype=np.uint8)

    raw_mask = cv2.resize(raw_mask, (w, h))
    mask = _clean_mask(raw_mask)
    is_valid, audit_status = _audit_mask_quality(mask, 1.0)
    final_mask = _center_and_pad_binary(mask)
    return is_valid, final_mask, audit_status, mask


# ---------------------------------------------------------------------------
# Classifier pipeline
# ---------------------------------------------------------------------------

def _mask_to_tensor(mask: np.ndarray) -> torch.Tensor:
    silhouette = Image.fromarray(mask).convert("L")
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
    ])
    return transform(silhouette).unsqueeze(0)


def _crop_constant_borders(img: np.ndarray, threshold: float = 15.0) -> np.ndarray:
    if img is None or img.size == 0:
        return img

    img_bgr = img[:, :, :3] if len(img.shape) == 3 and img.shape[2] == 4 else img
    if len(img_bgr.shape) == 3:
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    else:
        gray = img_bgr.copy()

    h, w = gray.shape

    row_means = np.mean(gray, axis=1)
    col_means = np.mean(gray, axis=0)

    left = 0
    while left < w and col_means[left] < threshold:
        left += 1

    right = w
    while right > left and col_means[right - 1] < threshold:
        right -= 1

    top = 0
    while top < h and row_means[top] < threshold:
        top += 1

    bottom = h
    while bottom > top and row_means[bottom - 1] < threshold:
        bottom -= 1

    if (right - left) >= 0.3 * w and (bottom - top) >= 0.3 * h:
        if left > 0 or right < w or top > 0 or bottom < h:
            return img[top:bottom, left:right]

    return img


def _classify_silhouette(
    final_mask: np.ndarray,
    model: PokedexNet,
    device: torch.device,
    id_to_name: Dict[int, str],
    detected_source: str,
    image_hash: str,
    include_debug: bool,
    original_width: int = 0,
    original_height: int = 0,
) -> PredictionResult:
    debug_b64 = None
    if include_debug:
        _, buffer = cv2.imencode(".png", final_mask)
        debug_b64 = base64.b64encode(buffer).decode("utf-8")

    input_tensor = _mask_to_tensor(final_mask).to(device)

    with torch.no_grad():
        output = model(input_tensor)
        probs = F.softmax(output, dim=1)
        top_probs, top_indices = torch.topk(probs, k=5, dim=1)

        top_5_list = [
            {
                "pokemon_id": int(top_indices[0][i].item()) + 1,
                "name":       id_to_name.get(int(top_indices[0][i].item()), "Unknown"),
                "confidence": float(top_probs[0][i].item()),
            }
            for i in range(5)
        ]

    best_pred = top_5_list[0]
    status = "CONFIDENT_POKEMON"
    top_class = int(top_indices[0][0].item())

    return PredictionResult(
        pokemon_id=int(best_pred["pokemon_id"]),
        name=str(best_pred["name"]),
        confidence=float(best_pred["confidence"]),
        detected_source=detected_source,
        top_5=top_5_list,
        debug_silhouette_b64=debug_b64,
        status=status,
        votes=[top_class],
        image_hash=image_hash,
        is_ood_screen=False,
        ensemble_metrics={
            "decision_tier": "Single_Model",
            "avg_max_logit": 0.0,
            "top1_probs": [],
            "avg_top1_prob": 0.0,
            "unique_predictions_count": 1,
        },
        original_width=original_width,
        original_height=original_height,
    )


# ---------------------------------------------------------------------------
# Main prediction entry point
# ---------------------------------------------------------------------------

def predict_from_bytes(
    file_bytes: bytes,
    models: List[PokedexNet],
    device: torch.device,
    id_to_name: Dict[int, str],
    include_debug: bool = False,
    progress: Callable[[Dict], None] | None = None,
) -> PredictionResult:
    image_hash = hashlib.sha256(file_bytes).hexdigest()

    np_arr = np.frombuffer(file_bytes, np.uint8)
    img    = cv2.imdecode(np_arr, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise ValueError("Could not decode image bytes. Unsupported or corrupted format.")

    original_height, original_width = img.shape[:2]
    max_dim = 8192
    if original_height > max_dim or original_width > max_dim:
        raise ValueError(
            f"Image dimensions ({original_width}x{original_height}) exceed maximum allowed ({max_dim}x{max_dim})"
        )

    img = _crop_constant_borders(img)

    detected_source = "birefnet"

    has_alpha = len(img.shape) == 3 and img.shape[2] == 4
    if len(img.shape) == 2:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    elif has_alpha:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
    else:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    h, w = img_rgb.shape[:2]
    resolutions = [settings.BIREFNET_FIRST_RESOLUTION, settings.BIREFNET_FALLBACK_RESOLUTION]

    if progress:
        progress({"phase": "segmenting"})

    best_failed_result = None

    for i, res in enumerate(resolutions):
        is_valid, final_mask, audit_status, raw_mask = _extract_birefnet_silhouette_for_res(img_rgb, res, w, h)

        if is_valid:
            result = _classify_silhouette(
                final_mask=final_mask,
                model=models[0],
                device=device,
                id_to_name=id_to_name,
                detected_source=detected_source,
                image_hash=image_hash,
                include_debug=include_debug,
                original_width=original_width,
                original_height=original_height,
            )
            return result
        else:
            debug_b64 = None
            if include_debug:
                _, buffer = cv2.imencode(".png", final_mask)
                debug_b64 = base64.b64encode(buffer).decode("utf-8")

            best_failed_result = PredictionResult(
                pokemon_id=0,
                name="Unknown",
                confidence=0.0,
                detected_source=detected_source,
                top_5=[],
                debug_silhouette_b64=debug_b64,
                status=audit_status,
                votes=[],
                image_hash=image_hash,
                is_ood_screen=True,
                ensemble_metrics={
                    "decision_tier": "Mask_Audit_Failure",
                    "avg_max_logit": 0.0,
                    "top1_probs": [],
                    "avg_top1_prob": 0.0,
                    "unique_predictions_count": 0,
                },
                original_width=original_width,
                original_height=original_height,
            )

        if i < len(resolutions) - 1:
            if progress:
                progress({"phase": "retrying"})

    return best_failed_result