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