modeling_dfine.py

"""
D-FINE HuggingFace Wrapper
Cho phép load model D-FINE thông qua transformers AutoModel.
Yêu cầu: Người dùng phải clone repo D-FINE và chạy code từ thư mục gốc của nó.
"""

import os
import sys
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image
from transformers import PretrainedConfig, PreTrainedModel, BaseImageProcessor


class DFINEConfig(PretrainedConfig):
    """Config class cho D-FINE model trên HuggingFace."""
    model_type = "dfine"

    def __init__(
        self,
        num_classes=4,
        resolution=640,
        checkpoint="best_stg2.pth",
        config_file="dfine_hgnetv2_x_custom.yml",
        **kwargs,
    ):
        self.num_classes = num_classes
        self.resolution = resolution
        self.checkpoint = checkpoint
        self.config_file = config_file
        super().__init__(**kwargs)


class DFINEImageProcessor(BaseImageProcessor):
    """Image processor: resize ảnh về kích thước chuẩn 640x640 và chuyển thành tensor."""

    def __init__(self, size=None, **kwargs):
        super().__init__(**kwargs)
        self.size = size or {"height": 640, "width": 640}

    def preprocess(self, images, **kwargs):
        if not isinstance(images, list):
            images = [images]
        processed = []
        for img in images:
            if not isinstance(img, Image.Image):
                img = Image.fromarray(img)
            img = img.convert("RGB")
            processed.append(img)
        return {"images": processed}

    def post_process_object_detection(self, outputs, threshold=0.5, target_sizes=None):
        """Lọc kết quả dự đoán theo ngưỡng confidence."""
        results = []
        for i, (boxes, scores, labels) in enumerate(
            zip(outputs["boxes"], outputs["scores"], outputs["labels"])
        ):
            keep = scores > threshold
            result = {
                "boxes": boxes[keep],
                "scores": scores[keep],
                "labels": labels[keep],
            }
            results.append(result)
        return results


class DFINEModel(PreTrainedModel):
    """
    HuggingFace wrapper cho D-FINE object detection model.

    Cách hoạt động:
    1. Tự động tải file checkpoint (.pth) từ HuggingFace Hub.
    2. Sử dụng source code D-FINE có sẵn trên máy để khởi tạo kiến trúc mạng.
    3. Load trọng số vào mạng và chạy inference.

    Yêu cầu: Người dùng phải clone repo D-FINE và chạy từ thư mục gốc.
    """

    config_class = DFINEConfig

    def __init__(self, config):
        super().__init__(config)
        self._model = None
        self._weights_path = None
        self._config_path = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        kwargs.pop("ignore_mismatched_sizes", None)
        config = kwargs.pop("config", None)
        if config is None:
            config = DFINEConfig.from_pretrained(
                pretrained_model_name_or_path,
                **{
                    k: v
                    for k, v in kwargs.items()
                    if k in ("trust_remote_code", "cache_dir", "revision", "token")
                },
            )
        model = cls(config)

        # Tải file checkpoint và config YAML từ HuggingFace Hub
        weights_file = config.checkpoint
        config_file = config.config_file

        local_weights = os.path.join(pretrained_model_name_or_path, weights_file)
        local_config = os.path.join(pretrained_model_name_or_path, config_file)

        if os.path.exists(local_weights):
            model._weights_path = local_weights
        else:
            from huggingface_hub import hf_hub_download

            model._weights_path = hf_hub_download(
                repo_id=pretrained_model_name_or_path, filename=weights_file
            )

        if os.path.exists(local_config):
            model._config_path = local_config
        else:
            from huggingface_hub import hf_hub_download

            model._config_path = hf_hub_download(
                repo_id=pretrained_model_name_or_path, filename=config_file
            )

        return model

    def _load_dfine(self, device=None):
        """Lazy load: Chỉ khởi tạo model khi thực sự cần chạy inference."""
        if self._model is not None:
            return

        # Thêm đường dẫn D-FINE vào sys.path để import được source code
        dfine_root = os.getcwd()
        if dfine_root not in sys.path:
            sys.path.insert(0, dfine_root)

        try:
            from src.core import YAMLConfig
        except ImportError:
            raise ImportError(
                "Không tìm thấy source code D-FINE! "
                "Hãy đảm bảo bạn đã clone repo D-FINE "
                "(git clone https://github.com/Peterande/D-FINE.git) "
                "và đang chạy code từ thư mục gốc của nó."
            )

        # Khởi tạo model từ YAML config
        cfg = YAMLConfig(self._config_path, resume=self._weights_path)

        if "HGNetv2" in cfg.yaml_cfg:
            cfg.yaml_cfg["HGNetv2"]["pretrained"] = False

        # Load trọng số
        checkpoint = torch.load(self._weights_path, map_location="cpu")
        if "ema" in checkpoint:
            state = checkpoint["ema"]["module"]
        else:
            state = checkpoint["model"]

        cfg.model.load_state_dict(state)

        # Chuyển sang deploy mode (loại bỏ các thành phần training)
        deploy_model = cfg.model.deploy()
        postprocessor = cfg.postprocessor.deploy()

        class _DeployModel(nn.Module):
            def __init__(self, model, postprocessor):
                super().__init__()
                self.model = model
                self.postprocessor = postprocessor

            def forward(self, images, orig_target_sizes):
                outputs = self.model(images)
                outputs = self.postprocessor(outputs, orig_target_sizes)
                return outputs

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        self._model = _DeployModel(deploy_model, postprocessor).to(device)
        self._model.eval()

    def forward(self, images, threshold=0.4):
        """
        Chạy inference trên danh sách ảnh PIL.

        Args:
            images: list[PIL.Image] hoặc một PIL.Image đơn lẻ.
            threshold: Ngưỡng confidence tối thiểu (mặc định 0.4).

        Returns:
            dict với keys: "boxes", "scores", "labels" (mỗi key là list of tensors).
        """
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self._load_dfine(device=device)

        if not isinstance(images, list):
            images = [images]

        transforms = T.Compose([T.Resize((640, 640)), T.ToTensor()])

        boxes_list, scores_list, labels_list = [], [], []
        for img in images:
            if not isinstance(img, Image.Image):
                img = Image.fromarray(img)
            img = img.convert("RGB")
            w, h = img.size
            orig_size = torch.tensor([[w, h]]).to(device)
            im_data = transforms(img).unsqueeze(0).to(device)

            with torch.no_grad():
                labels, boxes, scores = self._model(im_data, orig_size)

            boxes_list.append(boxes[0])
            scores_list.append(scores[0])
            labels_list.append(labels[0])

        return {"boxes": boxes_list, "scores": scores_list, "labels": labels_list}


# Đăng ký với HuggingFace AutoModel / AutoConfig
from transformers import AutoConfig, AutoModel

AutoConfig.register("dfine", DFINEConfig)
AutoModel.register(DFINEConfig, DFINEModel)