developer0hye/D-FINE
3
1"""2Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.3"""4import gradio as gr5import spaces6import os7import sys8import torch9import torch.nn as nn10import torchvision.transforms as T11import supervision as sv12from PIL import Image13import requests14import yaml15import numpy as np16import gc17 18from src.core import YAMLConfig19 20 21model_configs = {22 "dfine_n_coco":23 {"cfgfile": "configs/dfine/dfine_hgnetv2_n_coco.yml",24 "classinfofile": "configs/coco.yml",25 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_n_coco.pth"},26 "dfine_s_coco":27 {"cfgfile": "configs/dfine/dfine_hgnetv2_s_coco.yml",28 "classinfofile": "configs/coco.yml",29 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_coco.pth"},30 "dfine_m_coco":31 {"cfgfile": "configs/dfine/dfine_hgnetv2_m_coco.yml",32 "classinfofile": "configs/coco.yml",33 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_coco.pth"},34 "dfine_l_coco":35 {"cfgfile": "configs/dfine/dfine_hgnetv2_l_coco.yml",36 "classinfofile": "configs/coco.yml",37 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_coco.pth"},38 "dfine_x_coco":39 {"cfgfile": "configs/dfine/dfine_hgnetv2_x_coco.yml",40 "classinfofile": "configs/coco.yml",41 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_coco.pth"},42 "dfine_s_obj365":43 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_s_obj365.yml",44 "classinfofile": "configs/obj365.yml",45 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj365.pth"},46 "dfine_m_obj365":47 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_m_obj365.yml",48 "classinfofile": "configs/obj365.yml",49 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj365.pth"},50 "dfine_l_obj365":51 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml",52 "classinfofile": "configs/obj365.yml",53 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365.pth"},54 "dfine_l_obj365_e25":55 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml",56 "classinfofile": "configs/obj365.yml",57 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365_e25.pth"},58 "dfine_x_obj365":59 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_x_obj365.yml",60 "classinfofile": "configs/obj365.yml",61 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj365.pth"},62 "dfine_s_obj2coco":63 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_s_obj2coco.yml",64 "classinfofile": "configs/coco.yml",65 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj2coco.pth"},66 "dfine_m_obj2coco":67 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_m_obj2coco.yml",68 "classinfofile": "configs/coco.yml",69 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj2coco.pth"},70 "dfine_l_obj2coco_e25":71 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj2coco.yml",72 "classinfofile": "configs/coco.yml",73 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj2coco_e25.pth"},74 "dfine_x_obj2coco":75 {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_x_obj2coco.yml",76 "classinfofile": "configs/coco.yml",77 "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj2coco.pth"},78}79 80 81def download_weights(model_name):82 """Download model weights if not already present"""83 weights_url = model_configs[model_name]["weights"]84 # Directory path to save weight files85 weights_dir = os.path.join(os.path.dirname(__file__), "weights")86 # Weight file path87 weights_path = os.path.join(weights_dir, model_name + ".pth")88 89 # Create weights directory if it doesn't exist90 if not os.path.exists(weights_dir):91 os.makedirs(weights_dir)92 print(f"Created directory: {weights_dir}")93 94 # Check if file already exists95 if os.path.exists(weights_path):96 print(f"Weights file already exists at: {weights_path}")97 return weights_path98 99 # Download file100 print(f"Downloading weights from {weights_url} to {weights_path}...")101 102 response = requests.get(weights_url, stream=True)103 response.raise_for_status() # Check for download errors104 105 with open(weights_path, 'wb') as f:106 for chunk in response.iter_content(chunk_size=8192):107 f.write(chunk)108 109 print(f"Downloaded weights to: {weights_path}")110 return weights_path111 112@torch.no_grad()113def process_image_for_gradio(model, device, image, model_name, threshold=0.4):114 """Process image function for Gradio interface"""115 if isinstance(image, np.ndarray):116 # Convert NumPy array to PIL image117 im_pil = Image.fromarray(image)118 else:119 im_pil = image120 121 # Load class information122 classinfofile = model_configs[model_name]["classinfofile"]123 classinfo = yaml.load(open(classinfofile, "r"), Loader=yaml.FullLoader)["names"]124 indexing_method = "0-based" if "coco" in classinfofile else "1-based"125 126 w, h = im_pil.size127 orig_size = torch.tensor([[w, h]]).to(device)128 129 transforms = T.Compose(130 [131 T.Resize((640, 640)),132 T.ToTensor(),133 ]134 )135 im_data = transforms(im_pil).unsqueeze(0).to(device)136 137 output = model(im_data, orig_size)138 labels, boxes, scores = output139 140 # Visualize results141 detections = sv.Detections(142 xyxy=boxes[0].detach().cpu().numpy(),143 confidence=scores[0].detach().cpu().numpy(),144 class_id=labels[0].detach().cpu().numpy().astype(int),145 )146 detections = detections[detections.confidence > threshold]147 148 text_scale = sv.calculate_optimal_text_scale(resolution_wh=im_pil.size)149 line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=im_pil.size)150 151 box_annotator = sv.BoxAnnotator(thickness=line_thickness)152 label_annotator = sv.LabelAnnotator(text_scale=text_scale, smart_position=True)153 154 label_texts = [155 f"{classinfo[class_id if indexing_method == '0-based' else class_id - 1]} {confidence:.2f}"156 for class_id, confidence157 in zip(detections.class_id, detections.confidence)158 ]159 160 result_image = im_pil.copy()161 result_image = box_annotator.annotate(scene=result_image, detections=detections)162 result_image = label_annotator.annotate(163 scene=result_image,164 detections=detections,165 labels=label_texts166 )167 168 detection_info = [169 f"{classinfo[class_id if indexing_method == '0-based' else class_id - 1]}: {confidence:.2f}, bbox: [{xyxy[0]:.1f}, {xyxy[1]:.1f}, {xyxy[2]:.1f}, {xyxy[3]:.1f}]"170 for class_id, confidence, xyxy171 in zip(detections.class_id, detections.confidence, detections.xyxy)172 ]173 174 return result_image, "\n".join(detection_info)175 176 177class ModelWrapper(nn.Module):178 def __init__(self, cfg):179 super().__init__()180 self.model = cfg.model.deploy()181 self.postprocessor = cfg.postprocessor.deploy()182 183 def forward(self, images, orig_target_sizes):184 outputs = self.model(images)185 outputs = self.postprocessor(outputs, orig_target_sizes)186 return outputs187 188 189# YAMLConfig 클래스의 내부 상태를 초기화하는 함수 추가190def reset_yaml_config():191 """YAMLConfig 클래스의 내부 상태를 초기화"""192 # 클래스 내부에 캐싱된 정보가 있다면 삭제193 if hasattr(YAMLConfig, '_instances'):194 YAMLConfig._instances = {}195 if hasattr(YAMLConfig, '_configs'):196 YAMLConfig._configs = {}197 198 # 가능한 다른 모든 모듈 캐시 리셋199 import importlib200 for module_name in list(sys.modules.keys()):201 if module_name.startswith('src.'):202 try:203 importlib.reload(sys.modules[module_name])204 except:205 pass206 207def load_model(model_name):208 # 모델 로드 전에 CUDA 캐시와 가비지 컬렉션 정리209 if torch.cuda.is_available():210 torch.cuda.empty_cache()211 gc.collect()212 213 # YAMLConfig 내부 상태 초기화214 reset_yaml_config()215 216 cfgfile = model_configs[model_name]["cfgfile"]217 weights_path = download_weights(model_name)218 219 # 완전히 새로운 YAMLConfig 인스턴스 생성220 cfg = YAMLConfig(cfgfile, resume=weights_path)221 222 if "HGNetv2" in cfg.yaml_cfg:223 cfg.yaml_cfg["HGNetv2"]["pretrained"] = False224 225 checkpoint = torch.load(weights_path, map_location="cpu")226 state = checkpoint["ema"]["module"] if "ema" in checkpoint else checkpoint["model"]227 228 # 모델 생성 전 한번 더 확인229 torch.cuda.empty_cache()230 gc.collect()231 232 cfg.model.load_state_dict(state, strict=False)233 234 device = "cuda" if torch.cuda.is_available() else "cpu"235 model = ModelWrapper(cfg).to(device)236 model.eval()237 238 return model, device239 240@spaces.GPU241def process_image(image, model_name, confidence_threshold):242 """Main processing function for Gradio interface"""243 244 # 모든 사용 가능한 CUDA 장치 메모리 확보245 if torch.cuda.is_available():246 torch.cuda.empty_cache()247 248 # 모든 Python 객체 가비지 컬렉션249 gc.collect()250 251 try:252 print(f"Loading model: {model_name}")253 model, device = load_model(model_name)254 255 # 이미지 처리256 result = process_image_for_gradio(model, device, image, model_name, confidence_threshold)257 258 # 모델 객체 및 관련 데이터 명시적 제거259 del model260 261 finally:262 # 항상 메모리 정리 보장263 if torch.cuda.is_available():264 torch.cuda.empty_cache()265 gc.collect()266 267 return result268 269 270# Create Gradio interface271demo = gr.Interface(272 fn=process_image,273 inputs=[274 gr.Image(type="pil", label="Input Image"),275 gr.Dropdown(276 choices=list(model_configs.keys()), 277 value="dfine_n_coco", 278 label="Model Selection"279 ),280 gr.Slider(281 minimum=0.1, 282 maximum=0.9, 283 value=0.4, 284 step=0.05, 285 label="Confidence Threshold"286 )287 ],288 outputs=[289 gr.Image(type="pil", label="Detection Result"),290 gr.Textbox(label="Detected Objects")291 ],292 title="D-FINE Object Detection Demo",293 description="Upload an image to see object detection results using the D-FINE model. You can select different models and adjust the confidence threshold.",294 examples=[295 ["examples/image1.jpg", "dfine_n_coco", 0.4],296 ]297)298 299if __name__ == "__main__":300 # Launch the Gradio app301 demo.launch(share=True)