model.py

"""Car physical damage instance segmentation — Mask R-CNN R-101-DC5, pure PyTorch.

Toàn bộ kiến trúc nằm trong một file, không phụ thuộc Detectron2. Trọng số gốc
huấn luyện bằng Detectron2 0.6 nạp thẳng vào đây (tên module giữ nguyên 1:1).

    from cardamage import AutoModel
    model = AutoModel.from_pretrained("Naiscorp/car-damage-maskrcnn-r101-dc5").eval()
    overlay = model.inference(pil_image)

Kiến trúc:
    ResNet-101 DilatedC5 (res5 dilation=2, stride 16, 2048 ch, FrozenBN)
      -> RPN (StandardRPNHead, 15 anchor/vị trí)
      -> StandardROIHeads (ROIAlignV2 7x7 -> 2 FC 1024 -> cls/bbox)
      -> Mask head (ROIAlignV2 14x14 -> 4 conv 256 -> deconv -> 1x1 conv)
"""

from __future__ import annotations

import math
import os
import warnings
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont
from torch import nn
from torchvision.ops import batched_nms, roi_align

__all__ = [
    "CarDamageModel",
    "DAMAGE_CLASSES_VI",
    "DAMAGE_CLASSES_EN",
]

# Thứ tự nhãn khớp với chỉ số lớp của checkpoint (0..6).
DAMAGE_CLASSES_VI = [
    "Móp lõm",
    "Trầy sơn",
    "Rách",
    "Mất bộ phận",
    "Thủng",
    "Bể đèn",
    "Vỡ kính",
]
DAMAGE_CLASSES_EN = [
    "Dent",
    "Paint scratch",
    "Tear",
    "Missing part",
    "Puncture",
    "Broken lamp",
    "Broken glass",
]

# Bảng màu vẽ overlay, một màu cho mỗi lớp (RGB).
_PALETTE = [
    (255, 99, 71),
    (60, 179, 113),
    (65, 105, 225),
    (255, 165, 0),
    (186, 85, 211),
    (0, 206, 209),
    (255, 20, 147),
]


# --------------------------------------------------------------------------------------
# Layer cơ bản
# --------------------------------------------------------------------------------------
class FrozenBatchNorm2d(nn.Module):
    """BatchNorm với thống kê và tham số affine cố định (dùng khi fine-tune detection)."""

    def __init__(self, num_features: int, eps: float = 1e-5):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.register_buffer("weight", torch.ones(num_features))
        self.register_buffer("bias", torch.zeros(num_features))
        self.register_buffer("running_mean", torch.zeros(num_features))
        self.register_buffer("running_var", torch.ones(num_features) - eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.batch_norm(
            x,
            self.running_mean,
            self.running_var,
            self.weight,
            self.bias,
            training=False,
            eps=self.eps,
        )

    def __repr__(self) -> str:
        return f"FrozenBatchNorm2d(num_features={self.num_features}, eps={self.eps})"


class Conv2d(nn.Conv2d):
    """nn.Conv2d kèm norm/activation làm submodule — để tên key state_dict khớp checkpoint gốc."""

    def __init__(self, *args, norm: Optional[nn.Module] = None,
                 activation: Optional[nn.Module] = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.norm = norm
        self.activation = activation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding,
                     self.dilation, self.groups)
        if self.norm is not None:
            x = self.norm(x)
        if self.activation is not None:
            x = self.activation(x)
        return x


# --------------------------------------------------------------------------------------
# Backbone: ResNet-101 DilatedC5
# --------------------------------------------------------------------------------------
class BottleneckBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, *, bottleneck_channels: int,
                 stride: int = 1, num_groups: int = 1, stride_in_1x1: bool = False,
                 dilation: int = 1):
        super().__init__()
        self.stride = stride
        self.out_channels = out_channels

        if in_channels != out_channels:
            self.shortcut = Conv2d(in_channels, out_channels, kernel_size=1, stride=stride,
                                   bias=False, norm=FrozenBatchNorm2d(out_channels))
        else:
            self.shortcut = None

        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)

        self.conv1 = Conv2d(in_channels, bottleneck_channels, kernel_size=1, stride=stride_1x1,
                            bias=False, norm=FrozenBatchNorm2d(bottleneck_channels))
        self.conv2 = Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=3,
                            stride=stride_3x3, padding=1 * dilation, bias=False,
                            groups=num_groups, dilation=dilation,
                            norm=FrozenBatchNorm2d(bottleneck_channels))
        self.conv3 = Conv2d(bottleneck_channels, out_channels, kernel_size=1, bias=False,
                            norm=FrozenBatchNorm2d(out_channels))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.relu_(self.conv1(x))
        out = F.relu_(self.conv2(out))
        out = self.conv3(out)
        shortcut = x if self.shortcut is None else self.shortcut(x)
        out = out + shortcut
        return F.relu_(out)


class BasicStem(nn.Module):
    def __init__(self, in_channels: int = 3, out_channels: int = 64):
        super().__init__()
        self.conv1 = Conv2d(in_channels, out_channels, kernel_size=7, stride=2, padding=3,
                            bias=False, norm=FrozenBatchNorm2d(out_channels))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.relu_(self.conv1(x))
        return F.max_pool2d(x, kernel_size=3, stride=2, padding=1)


_NUM_BLOCKS_PER_STAGE = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}


class ResNetDC5(nn.Module):
    """ResNet chỉ trả về feature map res5, với res5 dùng dilation thay cho stride."""

    def __init__(self, depth: int = 101, res5_dilation: int = 2, num_groups: int = 1,
                 width_per_group: int = 64, stem_out_channels: int = 64,
                 res2_out_channels: int = 256, stride_in_1x1: bool = True):
        super().__init__()
        if depth not in _NUM_BLOCKS_PER_STAGE:
            raise ValueError(f"depth phải thuộc {sorted(_NUM_BLOCKS_PER_STAGE)}, nhận {depth}")
        if res5_dilation not in (1, 2):
            raise ValueError(f"res5_dilation phải là 1 hoặc 2, nhận {res5_dilation}")

        self.stem = BasicStem(3, stem_out_channels)

        num_blocks = _NUM_BLOCKS_PER_STAGE[depth]
        in_channels = stem_out_channels
        out_channels = res2_out_channels
        bottleneck_channels = num_groups * width_per_group
        current_stride = 4

        self.stage_names: List[str] = []
        for idx, stage_idx in enumerate(range(2, 6)):
            dilation = res5_dilation if stage_idx == 5 else 1
            first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
            blocks = []
            for i in range(num_blocks[idx]):
                blocks.append(
                    BottleneckBlock(
                        in_channels if i == 0 else out_channels,
                        out_channels,
                        bottleneck_channels=bottleneck_channels,
                        stride=first_stride if i == 0 else 1,
                        num_groups=num_groups,
                        stride_in_1x1=stride_in_1x1,
                        dilation=dilation,
                    )
                )
            name = f"res{stage_idx}"
            self.add_module(name, nn.Sequential(*blocks))
            self.stage_names.append(name)
            current_stride *= first_stride
            in_channels = out_channels
            out_channels *= 2
            bottleneck_channels *= 2

        self.out_channels = in_channels  # 2048 với R-50/101/152
        self.out_stride = current_stride  # 16 với DC5

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for name in self.stage_names:
            x = getattr(self, name)(x)
        return x


# --------------------------------------------------------------------------------------
# Anchor + box transform
# --------------------------------------------------------------------------------------
def _generate_cell_anchors(sizes: Sequence[float], aspect_ratios: Sequence[float]) -> torch.Tensor:
    anchors = []
    for size in sizes:
        area = size ** 2.0
        for aspect_ratio in aspect_ratios:
            w = math.sqrt(area / aspect_ratio)
            h = aspect_ratio * w
            anchors.append([-w / 2.0, -h / 2.0, w / 2.0, h / 2.0])
    return torch.tensor(anchors, dtype=torch.float32)


def _grid_anchors(cell_anchors: torch.Tensor, grid_h: int, grid_w: int, stride: int,
                  offset: float) -> torch.Tensor:
    device = cell_anchors.device
    shifts_x = torch.arange(offset * stride, grid_w * stride, step=stride,
                            dtype=torch.float32, device=device)
    shifts_y = torch.arange(offset * stride, grid_h * stride, step=stride,
                            dtype=torch.float32, device=device)
    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x, indexing="ij")
    shifts = torch.stack((shift_x.reshape(-1), shift_y.reshape(-1),
                          shift_x.reshape(-1), shift_y.reshape(-1)), dim=1)
    return (shifts.view(-1, 1, 4) + cell_anchors.view(1, -1, 4)).reshape(-1, 4)


_SCALE_CLAMP = math.log(1000.0 / 16)


def _apply_deltas(deltas: torch.Tensor, boxes: torch.Tensor,
                  weights: Tuple[float, float, float, float]) -> torch.Tensor:
    """Giải mã (dx, dy, dw, dh) thành toạ độ xyxy. Tương đương Box2BoxTransform.apply_deltas."""
    deltas = deltas.float()
    boxes = boxes.to(deltas.dtype)

    widths = boxes[:, 2] - boxes[:, 0]
    heights = boxes[:, 3] - boxes[:, 1]
    ctr_x = boxes[:, 0] + 0.5 * widths
    ctr_y = boxes[:, 1] + 0.5 * heights

    wx, wy, ww, wh = weights
    dx = deltas[:, 0::4] / wx
    dy = deltas[:, 1::4] / wy
    dw = torch.clamp(deltas[:, 2::4] / ww, max=_SCALE_CLAMP)
    dh = torch.clamp(deltas[:, 3::4] / wh, max=_SCALE_CLAMP)

    pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
    pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
    pred_w = torch.exp(dw) * widths[:, None]
    pred_h = torch.exp(dh) * heights[:, None]

    pred_boxes = torch.stack(
        (
            pred_ctr_x - 0.5 * pred_w,
            pred_ctr_y - 0.5 * pred_h,
            pred_ctr_x + 0.5 * pred_w,
            pred_ctr_y + 0.5 * pred_h,
        ),
        dim=-1,
    )
    return pred_boxes.reshape(deltas.shape)


def _clip_boxes(boxes: torch.Tensor, height, width) -> torch.Tensor:
    """Kẹp box vào khung ảnh. height/width nhận cả int lẫn Tensor 0-chiều.

    Nhánh Tensor tồn tại để khi trace sang ONNX, biên clip được suy ra động từ
    kích thước ảnh thay vì bị đóng băng thành hằng số của ảnh mẫu.
    """
    if torch.is_tensor(height) or torch.is_tensor(width):
        zero = boxes.new_zeros(())
        w = torch.as_tensor(width, dtype=boxes.dtype, device=boxes.device)
        h = torch.as_tensor(height, dtype=boxes.dtype, device=boxes.device)
        x1 = torch.minimum(torch.maximum(boxes[:, 0], zero), w)
        y1 = torch.minimum(torch.maximum(boxes[:, 1], zero), h)
        x2 = torch.minimum(torch.maximum(boxes[:, 2], zero), w)
        y2 = torch.minimum(torch.maximum(boxes[:, 3], zero), h)
    else:
        x1 = boxes[:, 0].clamp(min=0, max=width)
        y1 = boxes[:, 1].clamp(min=0, max=height)
        x2 = boxes[:, 2].clamp(min=0, max=width)
        y2 = boxes[:, 3].clamp(min=0, max=height)
    return torch.stack((x1, y1, x2, y2), dim=-1)


def _nonempty(boxes: torch.Tensor, threshold: float = 0.0) -> torch.Tensor:
    widths = boxes[:, 2] - boxes[:, 0]
    heights = boxes[:, 3] - boxes[:, 1]
    return (widths > threshold) & (heights > threshold)


# --------------------------------------------------------------------------------------
# RPN
# --------------------------------------------------------------------------------------
class StandardRPNHead(nn.Module):
    def __init__(self, in_channels: int, num_anchors: int):
        super().__init__()
        self.conv = Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1,
                           activation=nn.ReLU())
        self.objectness_logits = nn.Conv2d(in_channels, num_anchors, kernel_size=1, stride=1)
        self.anchor_deltas = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=1, stride=1)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        t = self.conv(x)
        return self.objectness_logits(t), self.anchor_deltas(t)


class RPN(nn.Module):
    def __init__(self, in_channels: int, anchor_sizes: Sequence[float],
                 aspect_ratios: Sequence[float], stride: int, offset: float = 0.0,
                 pre_nms_topk: int = 6000, post_nms_topk: int = 1000,
                 nms_thresh: float = 0.7, min_box_size: float = 0.0,
                 box_reg_weights: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0)):
        super().__init__()
        num_anchors = len(anchor_sizes) * len(aspect_ratios)
        self.rpn_head = StandardRPNHead(in_channels, num_anchors)
        self.register_buffer(
            "cell_anchors", _generate_cell_anchors(anchor_sizes, aspect_ratios), persistent=False
        )
        self.stride = stride
        self.offset = offset
        self.pre_nms_topk = pre_nms_topk
        self.post_nms_topk = post_nms_topk
        self.nms_thresh = nms_thresh
        self.min_box_size = min_box_size
        self.box_reg_weights = tuple(box_reg_weights)

    def forward(self, feature: torch.Tensor,
                image_sizes: List[Tuple[int, int]]) -> List[torch.Tensor]:
        n, _, grid_h, grid_w = feature.shape
        anchors = _grid_anchors(self.cell_anchors, grid_h, grid_w, self.stride, self.offset)

        logits, deltas = self.rpn_head(feature)
        # (N, A, H, W) -> (N, H*W*A)
        logits = logits.permute(0, 2, 3, 1).flatten(1)
        # (N, A*4, H, W) -> (N, H*W*A, 4)
        deltas = (
            deltas.view(n, -1, 4, grid_h, grid_w).permute(0, 3, 4, 1, 2).flatten(1, -2)
        )

        proposals = _apply_deltas(
            deltas.reshape(-1, 4),
            anchors.unsqueeze(0).expand(n, -1, -1).reshape(-1, 4),
            self.box_reg_weights,
        ).view(n, -1, 4)

        return self._select_proposals(proposals, logits, image_sizes)

    def _select_proposals(self, proposals: torch.Tensor, logits: torch.Tensor,
                          image_sizes: List[Tuple[int, int]]) -> List[torch.Tensor]:
        num_proposals = min(logits.shape[1], self.pre_nms_topk)
        topk_scores, topk_idx = logits.topk(num_proposals, dim=1)
        batch_idx = torch.arange(proposals.shape[0], device=proposals.device)
        topk_proposals = proposals[batch_idx[:, None], topk_idx]

        results = []
        for i, (height, width) in enumerate(image_sizes):
            boxes = topk_proposals[i]
            scores = topk_scores[i]

            valid = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores)
            if not valid.all():
                boxes, scores = boxes[valid], scores[valid]

            boxes = _clip_boxes(boxes, height, width)
            keep = _nonempty(boxes, self.min_box_size)
            boxes, scores = boxes[keep], scores[keep]

            # Một feature level duy nhất -> mọi box cùng "lớp" khi batched_nms.
            lvl = torch.zeros_like(scores, dtype=torch.int64)
            keep = batched_nms(boxes.float(), scores, lvl, self.nms_thresh)[: self.post_nms_topk]
            results.append(boxes[keep])
        return results


# --------------------------------------------------------------------------------------
# ROI heads
# --------------------------------------------------------------------------------------
def _pooler_format(boxes_per_image: List[torch.Tensor]) -> torch.Tensor:
    out = []
    for i, boxes in enumerate(boxes_per_image):
        idx = torch.full((boxes.shape[0], 1), i, dtype=boxes.dtype, device=boxes.device)
        out.append(torch.cat((idx, boxes), dim=1))
    return torch.cat(out, dim=0)


def _roi_align(feature: torch.Tensor, boxes_per_image: List[torch.Tensor], output_size: int,
               spatial_scale: float) -> torch.Tensor:
    """ROIAlignV2 (aligned=True, sampling_ratio=0) — khớp mặc định của Detectron2."""
    rois = _pooler_format(boxes_per_image)
    if rois.shape[0] == 0:
        return feature.new_zeros((0, feature.shape[1], output_size, output_size))
    return roi_align(
        feature,
        rois.to(dtype=feature.dtype),
        output_size=(output_size, output_size),
        spatial_scale=spatial_scale,
        sampling_ratio=0,
        aligned=True,
    )


class FastRCNNConvFCHead(nn.Module):
    def __init__(self, in_channels: int, pooler_resolution: int, fc_dim: int = 1024,
                 num_fc: int = 2):
        super().__init__()
        self.num_fc = num_fc
        dim = in_channels * pooler_resolution * pooler_resolution
        for k in range(num_fc):
            self.add_module(f"fc{k + 1}", nn.Linear(dim, fc_dim))
            dim = fc_dim
        self.out_channels = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.flatten(x, start_dim=1)
        for k in range(self.num_fc):
            x = F.relu(getattr(self, f"fc{k + 1}")(x))
        return x


class FastRCNNOutputLayers(nn.Module):
    def __init__(self, input_size: int, num_classes: int):
        super().__init__()
        self.cls_score = nn.Linear(input_size, num_classes + 1)
        self.bbox_pred = nn.Linear(input_size, num_classes * 4)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.cls_score(x), self.bbox_pred(x)


class MaskRCNNConvUpsampleHead(nn.Module):
    def __init__(self, in_channels: int, num_classes: int, conv_dim: int = 256,
                 num_conv: int = 4):
        super().__init__()
        self.num_conv = num_conv
        cur = in_channels
        for k in range(num_conv):
            self.add_module(
                f"mask_fcn{k + 1}",
                Conv2d(cur, conv_dim, kernel_size=3, stride=1, padding=1, bias=True),
            )
            cur = conv_dim
        self.deconv = nn.ConvTranspose2d(cur, conv_dim, kernel_size=2, stride=2, padding=0)
        self.predictor = Conv2d(conv_dim, num_classes, kernel_size=1, stride=1, padding=0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for k in range(self.num_conv):
            x = F.relu(getattr(self, f"mask_fcn{k + 1}")(x))
        x = F.relu(self.deconv(x))
        return self.predictor(x)


class StandardROIHeads(nn.Module):
    def __init__(self, in_channels: int, num_classes: int, feature_stride: int,
                 box_pooler_resolution: int = 7, mask_pooler_resolution: int = 14,
                 fc_dim: int = 1024, mask_conv_dim: int = 256, mask_num_conv: int = 4,
                 box_reg_weights: Tuple[float, float, float, float] = (10.0, 10.0, 5.0, 5.0),
                 score_thresh: float = 0.7, nms_thresh: float = 0.5,
                 detections_per_image: int = 100):
        super().__init__()
        self.num_classes = num_classes
        self.spatial_scale = 1.0 / feature_stride
        self.box_pooler_resolution = box_pooler_resolution
        self.mask_pooler_resolution = mask_pooler_resolution
        self.box_reg_weights = tuple(box_reg_weights)
        self.score_thresh = score_thresh
        self.nms_thresh = nms_thresh
        self.detections_per_image = detections_per_image

        self.box_head = FastRCNNConvFCHead(in_channels, box_pooler_resolution, fc_dim=fc_dim)
        self.box_predictor = FastRCNNOutputLayers(self.box_head.out_channels, num_classes)
        self.mask_head = MaskRCNNConvUpsampleHead(in_channels, num_classes,
                                                  conv_dim=mask_conv_dim, num_conv=mask_num_conv)

    def forward(self, feature: torch.Tensor, proposals: List[torch.Tensor],
                image_sizes: List[Tuple[int, int]]) -> List[Dict[str, torch.Tensor]]:
        box_features = _roi_align(feature, proposals, self.box_pooler_resolution,
                                  self.spatial_scale)
        box_features = self.box_head(box_features)
        scores, proposal_deltas = self.box_predictor(box_features)

        num_per_image = [p.shape[0] for p in proposals]
        all_proposals = torch.cat(proposals, dim=0) if len(proposals) else \
            feature.new_zeros((0, 4))
        pred_boxes = _apply_deltas(proposal_deltas, all_proposals, self.box_reg_weights)
        probs = F.softmax(scores, dim=-1)

        results: List[Dict[str, torch.Tensor]] = []
        for boxes_i, scores_i, image_size in zip(
            pred_boxes.split(num_per_image), probs.split(num_per_image, dim=0), image_sizes
        ):
            results.append(self._inference_single_image(boxes_i, scores_i, image_size))

        boxes_for_mask = [r["boxes"] for r in results]
        if sum(b.shape[0] for b in boxes_for_mask) > 0:
            mask_features = _roi_align(feature, boxes_for_mask, self.mask_pooler_resolution,
                                       self.spatial_scale)
            mask_logits = self.mask_head(mask_features)
            classes = torch.cat([r["classes"] for r in results], dim=0)
            idx = torch.arange(mask_logits.shape[0], device=mask_logits.device)
            mask_probs = mask_logits[idx, classes][:, None].sigmoid()
            for r, m in zip(results, mask_probs.split([b.shape[0] for b in boxes_for_mask])):
                r["masks"] = m
        else:
            for r in results:
                r["masks"] = feature.new_zeros(
                    (0, 1, self.mask_pooler_resolution * 2, self.mask_pooler_resolution * 2)
                )
        return results

    def _inference_single_image(self, boxes: torch.Tensor, scores: torch.Tensor,
                                image_size: Tuple[int, int]) -> Dict[str, torch.Tensor]:
        height, width = image_size
        valid = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores).all(dim=1)
        if not valid.all():
            boxes, scores = boxes[valid], scores[valid]

        scores = scores[:, :-1]  # bỏ cột background
        num_bbox_reg_classes = boxes.shape[1] // 4
        boxes = _clip_boxes(boxes.reshape(-1, 4), height, width)
        boxes = boxes.view(-1, num_bbox_reg_classes, 4)

        filter_mask = scores > self.score_thresh
        filter_inds = filter_mask.nonzero()
        if num_bbox_reg_classes == 1:
            boxes = boxes[filter_inds[:, 0], 0]
        else:
            boxes = boxes[filter_mask]
        scores = scores[filter_mask]

        keep = batched_nms(boxes.float(), scores, filter_inds[:, 1], self.nms_thresh)
        if self.detections_per_image >= 0:
            keep = keep[: self.detections_per_image]
        return {
            "boxes": boxes[keep],
            "scores": scores[keep],
            "classes": filter_inds[keep][:, 1],
        }


# --------------------------------------------------------------------------------------
# Dán mask về ảnh gốc
# --------------------------------------------------------------------------------------
_GPU_MEM_LIMIT = 1024 ** 3
_BYTES_PER_FLOAT = 4


def _do_paste_mask(masks: torch.Tensor, boxes: torch.Tensor, img_h: int, img_w: int,
                   skip_empty: bool = True):
    device = masks.device
    if skip_empty:
        x0_int, y0_int = torch.clamp(
            boxes.min(dim=0).values.floor()[:2] - 1, min=0
        ).to(dtype=torch.int32)
        x1_int = torch.clamp(boxes[:, 2].max().ceil() + 1, max=img_w).to(dtype=torch.int32)
        y1_int = torch.clamp(boxes[:, 3].max().ceil() + 1, max=img_h).to(dtype=torch.int32)
    else:
        x0_int, y0_int = 0, 0
        x1_int, y1_int = img_w, img_h

    x0, y0, x1, y1 = torch.split(boxes, 1, dim=1)
    n = masks.shape[0]

    img_y = torch.arange(y0_int, y1_int, device=device, dtype=torch.float32) + 0.5
    img_x = torch.arange(x0_int, x1_int, device=device, dtype=torch.float32) + 0.5
    img_y = (img_y - y0) / (y1 - y0) * 2 - 1
    img_x = (img_x - x0) / (x1 - x0) * 2 - 1

    gx = img_x[:, None, :].expand(n, img_y.size(1), img_x.size(1))
    gy = img_y[:, :, None].expand(n, img_y.size(1), img_x.size(1))
    grid = torch.stack([gx, gy], dim=3)

    img_masks = F.grid_sample(masks.to(dtype=torch.float32), grid, align_corners=False)
    if skip_empty:
        return img_masks[:, 0], (slice(int(y0_int), int(y1_int)), slice(int(x0_int), int(x1_int)))
    return img_masks[:, 0], ()


def paste_masks_in_image(masks: torch.Tensor, boxes: torch.Tensor,
                         image_shape: Tuple[int, int], threshold: float = 0.5) -> torch.Tensor:
    """masks: (N, Hm, Wm) xác suất; boxes: (N, 4) trên toạ độ ảnh gốc."""
    n = len(masks)
    if n == 0:
        return torch.empty((0,) + tuple(image_shape), dtype=torch.bool, device=masks.device)

    img_h, img_w = image_shape
    device = boxes.device
    if device.type == "cpu":
        num_chunks = n
    else:
        num_chunks = int(
            np.ceil(n * int(img_h) * int(img_w) * _BYTES_PER_FLOAT / _GPU_MEM_LIMIT)
        )
        num_chunks = max(1, min(num_chunks, n))
    chunks = torch.chunk(torch.arange(n, device=device), num_chunks)

    img_masks = torch.zeros(n, img_h, img_w, device=device,
                            dtype=torch.bool if threshold >= 0 else torch.uint8)
    for inds in chunks:
        masks_chunk, spatial_inds = _do_paste_mask(
            masks[inds, None, :, :], boxes[inds], img_h, img_w,
            skip_empty=device.type == "cpu",
        )
        if threshold >= 0:
            masks_chunk = (masks_chunk >= threshold).to(dtype=torch.bool)
        else:
            masks_chunk = (masks_chunk * 255).to(dtype=torch.uint8)
        img_masks[(inds,) + spatial_inds] = masks_chunk
    return img_masks


# --------------------------------------------------------------------------------------
# Tiện ích ảnh
# --------------------------------------------------------------------------------------
def _resize_shortest_edge(height: int, width: int, short_edge: int,
                          max_size: int) -> Tuple[int, int]:
    scale = short_edge * 1.0 / min(height, width)
    if height < width:
        new_h, new_w = short_edge, scale * width
    else:
        new_h, new_w = scale * height, short_edge
    if max(new_h, new_w) > max_size:
        scale = max_size * 1.0 / max(new_h, new_w)
        new_h *= scale
        new_w *= scale
    return int(new_h + 0.5), int(new_w + 0.5)


_FONT_CACHE: Dict[Tuple[int, str], Tuple[ImageFont.ImageFont, bool]] = {}


def _font_paths() -> List[str]:
    """Font ứng viên, xếp theo thứ tự ưu tiên."""
    paths: List[str] = []
    # matplotlib luôn đóng gói kèm DejaVuSans (phủ đủ dấu tiếng Việt) và có sẵn trên
    # Kaggle, Colab, hầu hết môi trường khoa học dữ liệu — kể cả khi máy không cài font nào.
    try:
        import matplotlib

        paths.append(os.path.join(matplotlib.get_data_path(), "fonts", "ttf", "DejaVuSans.ttf"))
    except Exception:  # matplotlib không bắt buộc
        pass
    paths += [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/TTF/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
        "/usr/share/fonts/truetype/noto/NotoSans-Regular.ttf",
        "/usr/share/fonts/truetype/freefont/FreeSans.ttf",
        "/Library/Fonts/Arial Unicode.ttf",
        "/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
        "/System/Library/Fonts/Supplemental/Arial.ttf",
        "C:\\Windows\\Fonts\\segoeui.ttf",
        "C:\\Windows\\Fonts\\arial.ttf",
    ]
    return paths


def _missing_glyphs(font: ImageFont.ImageFont, text: str) -> List[str]:
    """Ký tự trong ``text`` mà ``font`` không có glyph.

    Không dùng thư viện ngoài: ký tự thiếu glyph được FreeType vẽ thành ``.notdef``,
    giống hệt một codepoint chắc chắn không tồn tại (vùng Private Use U+E000).
    """
    try:
        tofu = font.getmask("\ue000", mode="L")
        ref, ref_size = bytes(tofu), tofu.size
    except Exception:
        return []  # không kiểm tra được thì coi như đủ, thà vẽ còn hơn không
    out: List[str] = []
    for ch in dict.fromkeys(text):
        if ch.isspace():
            continue
        try:
            mask = font.getmask(ch, mode="L")
        except Exception:
            continue
        if mask.size == ref_size and bytes(mask) == ref:
            out.append(ch)
    return out


def _load_font(size: int, text: str = "") -> Tuple[ImageFont.ImageFont, bool]:
    """Trả về ``(font, đủ_glyph_cho_text)``.

    Mở được file font là CHƯA ĐỦ: font mặc định của Pillow chỉ có ~133 codepoint nên
    mọi dấu tiếng Việt bị vẽ thành ô vuông. Vì vậy hàm này kiểm tra glyph thật sự và
    đi tiếp sang ứng viên khác nếu thiếu.
    """
    key = (size, text)
    cached = _FONT_CACHE.get(key)
    if cached is not None:
        return cached

    fallback: Optional[ImageFont.ImageFont] = None
    for path in _font_paths():
        try:
            font = ImageFont.truetype(path, size)
        except OSError:
            continue
        if not text or not _missing_glyphs(font, text):
            _FONT_CACHE[key] = (font, True)
            return _FONT_CACHE[key]
        if fallback is None:
            fallback = font

    if fallback is None:
        try:  # Pillow >= 9.2 cho phép đặt cỡ cho font mặc định
            fallback = ImageFont.load_default(size=size)
        except TypeError:
            fallback = ImageFont.load_default()

    result = (fallback, not text or not _missing_glyphs(fallback, text))
    _FONT_CACHE[key] = result
    return result


# --------------------------------------------------------------------------------------
# Model chính
# --------------------------------------------------------------------------------------
try:
    from huggingface_hub import PyTorchModelHubMixin

    _MIXIN_BASES = (nn.Module, PyTorchModelHubMixin)
    _MIXIN_KWARGS = dict(
        pipeline_tag="image-segmentation",
        license="other",
        tags=["car-damage", "instance-segmentation", "mask-rcnn", "insurance", "detectron2"],
    )
except ImportError:  # pragma: no cover - chỉ xảy ra khi thiếu huggingface_hub
    _MIXIN_BASES = (nn.Module,)
    _MIXIN_KWARGS = {}


class CarDamageModel(*_MIXIN_BASES, **_MIXIN_KWARGS):  # type: ignore[misc]
    """Mask R-CNN R-101-DC5 phát hiện 7 loại tổn thất vật lý trên ô tô.

    Mọi tham số của ``__init__`` đều JSON-serializable nên
    ``save_pretrained``/``push_to_hub`` tự sinh được ``config.json``.
    """

    def __init__(
        self,
        num_classes: int = 7,
        class_names: List[str] = DAMAGE_CLASSES_VI,
        class_names_en: List[str] = DAMAGE_CLASSES_EN,
        backbone_depth: int = 101,
        res5_dilation: int = 2,
        stride_in_1x1: bool = True,
        anchor_sizes: List[float] = [32.0, 64.0, 128.0, 256.0, 512.0],
        anchor_aspect_ratios: List[float] = [0.5, 1.0, 2.0],
        anchor_offset: float = 0.0,
        rpn_pre_nms_topk: int = 6000,
        rpn_post_nms_topk: int = 1000,
        rpn_nms_thresh: float = 0.7,
        rpn_min_box_size: float = 0.0,
        rpn_box_reg_weights: List[float] = [1.0, 1.0, 1.0, 1.0],
        box_pooler_resolution: int = 7,
        mask_pooler_resolution: int = 14,
        box_head_fc_dim: int = 1024,
        mask_head_conv_dim: int = 256,
        mask_head_num_conv: int = 4,
        box_reg_weights: List[float] = [10.0, 10.0, 5.0, 5.0],
        score_thresh: float = 0.7,
        nms_thresh: float = 0.5,
        detections_per_image: int = 100,
        mask_threshold: float = 0.5,
        min_size_test: int = 800,
        max_size_test: int = 1333,
        pixel_mean: List[float] = [103.53, 116.28, 123.675],
        pixel_std: List[float] = [1.0, 1.0, 1.0],
        input_format: str = "BGR",
    ):
        # Các default dạng list ở trên là cố ý: PyTorchModelHubMixin ghi thẳng giá trị
        # nhận được vào config.json, nên dùng sentinel None sẽ làm mất thông tin nhãn.
        # Chúng không bao giờ bị sửa tại chỗ — mọi chỗ dùng đều copy bằng list(...).
        super().__init__()

        self.num_classes = num_classes
        self.class_names = list(class_names)
        self.class_names_en = list(class_names_en)
        if len(self.class_names) != num_classes:
            raise ValueError(
                f"class_names có {len(self.class_names)} phần tử nhưng num_classes={num_classes}"
            )
        if len(self.class_names_en) != num_classes:
            raise ValueError(
                f"class_names_en có {len(self.class_names_en)} phần tử "
                f"nhưng num_classes={num_classes}"
            )
        if input_format not in ("BGR", "RGB"):
            raise ValueError(f"input_format phải là 'BGR' hoặc 'RGB', nhận {input_format!r}")

        self.mask_threshold = mask_threshold
        self.min_size_test = min_size_test
        self.max_size_test = max_size_test
        self.input_format = input_format

        self.backbone = ResNetDC5(
            depth=backbone_depth, res5_dilation=res5_dilation, stride_in_1x1=stride_in_1x1
        )
        self.proposal_generator = RPN(
            in_channels=self.backbone.out_channels,
            anchor_sizes=anchor_sizes,
            aspect_ratios=anchor_aspect_ratios,
            stride=self.backbone.out_stride,
            offset=anchor_offset,
            pre_nms_topk=rpn_pre_nms_topk,
            post_nms_topk=rpn_post_nms_topk,
            nms_thresh=rpn_nms_thresh,
            min_box_size=rpn_min_box_size,
            box_reg_weights=tuple(rpn_box_reg_weights),
        )
        self.roi_heads = StandardROIHeads(
            in_channels=self.backbone.out_channels,
            num_classes=num_classes,
            feature_stride=self.backbone.out_stride,
            box_pooler_resolution=box_pooler_resolution,
            mask_pooler_resolution=mask_pooler_resolution,
            fc_dim=box_head_fc_dim,
            mask_conv_dim=mask_head_conv_dim,
            mask_num_conv=mask_head_num_conv,
            box_reg_weights=tuple(box_reg_weights),
            score_thresh=score_thresh,
            nms_thresh=nms_thresh,
            detections_per_image=detections_per_image,
        )

        self.register_buffer(
            "pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), persistent=False
        )
        self.register_buffer(
            "pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), persistent=False
        )

    # -- thuộc tính tiện dụng -------------------------------------------------------
    @property
    def device(self) -> torch.device:
        return next(self.parameters()).device

    # -- ngưỡng chỉnh được lúc chạy --------------------------------------------------
    @property
    def score_thresh(self) -> float:
        return self.roi_heads.score_thresh

    @score_thresh.setter
    def score_thresh(self, value: float) -> None:
        self.roi_heads.score_thresh = float(value)

    # -- forward thô -----------------------------------------------------------------
    def forward(self, images: torch.Tensor,
                image_sizes: Optional[List[Tuple[int, int]]] = None
                ) -> List[Dict[str, torch.Tensor]]:
        """images: (N, 3, H, W) đã chuẩn hoá (trừ mean, chia std), đúng ``input_format``.

        Trả về list dict theo từng ảnh, toạ độ nằm trên hệ ảnh **đã resize**.
        """
        if images.dim() != 4:
            raise ValueError(f"images phải có shape (N,3,H,W), nhận {tuple(images.shape)}")
        if image_sizes is None:
            image_sizes = [(images.shape[2], images.shape[3])] * images.shape[0]

        features = self.backbone(images)
        proposals = self.proposal_generator(features, image_sizes)
        return self.roi_heads(features, proposals, image_sizes)

    # -- tiền xử lý ------------------------------------------------------------------
    def _preprocess(self, image: Image.Image) -> Tuple[torch.Tensor, Tuple[int, int]]:
        rgb = image.convert("RGB")
        orig_w, orig_h = rgb.size
        new_h, new_w = _resize_shortest_edge(orig_h, orig_w, self.min_size_test,
                                             self.max_size_test)
        resized = rgb.resize((new_w, new_h), Image.BILINEAR)

        array = np.asarray(resized, dtype=np.uint8)
        if self.input_format == "BGR":
            array = array[:, :, ::-1]
        tensor = torch.as_tensor(np.ascontiguousarray(array.transpose(2, 0, 1)),
                                 dtype=torch.float32)
        tensor = (tensor - self.pixel_mean.cpu()) / self.pixel_std.cpu()
        return tensor.unsqueeze(0), (orig_h, orig_w)

    # -- suy luận có cấu trúc --------------------------------------------------------
    @torch.no_grad()
    def predict(self, image: Image.Image) -> Dict[str, object]:
        """Nhận PIL.Image, trả về dict: boxes/scores/classes/labels/masks trên ảnh gốc."""
        self.eval()
        device = self.device
        tensor, (orig_h, orig_w) = self._preprocess(image)
        tensor = tensor.to(device)
        resized_h, resized_w = tensor.shape[2], tensor.shape[3]

        out = self.forward(tensor, [(resized_h, resized_w)])[0]

        boxes = out["boxes"]
        scores = out["scores"]
        classes = out["classes"]
        masks = out["masks"]

        scale_x = orig_w / resized_w
        scale_y = orig_h / resized_h
        if boxes.shape[0]:
            boxes = boxes.clone()
            boxes[:, 0::2] *= scale_x
            boxes[:, 1::2] *= scale_y
            boxes = _clip_boxes(boxes, orig_h, orig_w)
            keep = _nonempty(boxes)
            boxes, scores, classes, masks = boxes[keep], scores[keep], classes[keep], masks[keep]

        pasted = paste_masks_in_image(
            masks[:, 0, :, :] if masks.shape[0] else masks.new_zeros((0, 1, 1)),
            boxes,
            (orig_h, orig_w),
            self.mask_threshold,
        )

        classes_np = classes.cpu().numpy()
        return {
            "boxes": boxes.cpu().numpy(),
            "scores": scores.cpu().numpy(),
            "classes": classes_np,
            "labels": [self.class_names[int(c)] for c in classes_np],
            "labels_en": [self.class_names_en[int(c)] for c in classes_np],
            "masks": pasted.cpu().numpy(),
            "image_size": (orig_h, orig_w),
        }

    # -- suy luận trả ảnh ------------------------------------------------------------
    @torch.no_grad()
    def inference(self, image: Image.Image, alpha: float = 0.45, draw_boxes: bool = True,
                  draw_labels: bool = True, language: str = "vi") -> Image.Image:
        """Nhận PIL.Image, trả về PIL.Image RGB đã vẽ mask + box + nhãn tổn thất."""
        result = self.predict(image)
        return self.visualize(image, result, alpha=alpha, draw_boxes=draw_boxes,
                              draw_labels=draw_labels, language=language)

    def visualize(self, image: Image.Image, result: Dict[str, object], alpha: float = 0.45,
                  draw_boxes: bool = True, draw_labels: bool = True,
                  language: str = "vi") -> Image.Image:
        """Vẽ kết quả của :meth:`predict` lên ảnh gốc."""
        base = image.convert("RGB")
        canvas = np.asarray(base, dtype=np.float32).copy()

        masks = result["masks"]
        classes = result["classes"]
        for mask, cls in zip(masks, classes):
            color = np.array(_PALETTE[int(cls) % len(_PALETTE)], dtype=np.float32)
            sel = mask.astype(bool)
            if sel.any():
                canvas[sel] = canvas[sel] * (1.0 - alpha) + color * alpha

        out = Image.fromarray(canvas.clip(0, 255).astype(np.uint8))
        if not (draw_boxes or draw_labels):
            return out

        draw = ImageDraw.Draw(out)
        font_size = max(14, int(round(min(out.size) * 0.028)))
        names = result["labels_en"] if language == "en" else result["labels"]

        # Chọn font theo ĐÚNG chuỗi sắp vẽ. Nếu máy không có font nào phủ được dấu
        # tiếng Việt thì hạ xuống nhãn tiếng Anh (thuần ASCII) — chữ đọc được vẫn hơn
        # một hàng ô vuông.
        probe = "".join(names) + "0123456789. "
        font, ok = _load_font(font_size, probe)
        if not ok and language != "en":
            warnings.warn(
                "Khong tim thay font nao ve duoc dau tieng Viet, dang dung nhan tieng Anh. "
                "Cai mot font Unicode (vi du: apt-get install fonts-dejavu-core, "
                "hoac pip install matplotlib) roi chay lai de co nhan tieng Viet.",
                RuntimeWarning, stacklevel=2,
            )
            names = result["labels_en"]
            font, _ = _load_font(font_size, "".join(names) + "0123456789. ")

        for box, score, cls, name in zip(result["boxes"], result["scores"], classes, names):
            color = _PALETTE[int(cls) % len(_PALETTE)]
            x0, y0, x1, y1 = [float(v) for v in box]
            if draw_boxes:
                draw.rectangle([x0, y0, x1, y1], outline=color, width=max(2, font_size // 7))
            if draw_labels:
                text = f"{name} {score:.2f}"
                try:
                    tx0, ty0, tx1, ty1 = draw.textbbox((0, 0), text, font=font)
                    tw, th = tx1 - tx0, ty1 - ty0
                except AttributeError:  # Pillow rất cũ
                    tw, th = draw.textsize(text, font=font)
                ty = max(0.0, y0 - th - 4)
                draw.rectangle([x0, ty, x0 + tw + 6, ty + th + 4], fill=color)
                draw.text((x0 + 3, ty + 2), text, fill=(255, 255, 255), font=font)
        return out

    # -- nạp trọng số Detectron2 gốc --------------------------------------------------
    @torch.no_grad()
    def load_detectron2_checkpoint(self, path: str, strict: bool = True):
        """Nạp file .pth do Detectron2 sinh ra (kể cả bản còn optimizer state)."""
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        state = ckpt.get("model", ckpt) if isinstance(ckpt, dict) else ckpt
        state = {k: v for k, v in state.items()
                 if not k.endswith(("pixel_mean", "pixel_std"))}
        return self.load_state_dict(state, strict=strict)