"""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)