Spaces:
Paused
Paused
File size: 3,737 Bytes
8bb7465 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import logging
from pathlib import Path
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from .prediction_helper import ResnetCarDamagePredictor
logger = logging.getLogger(__name__)
MODEL_CONFIG = {
"resnet_onnx": {
"repo_id": "junaid17/car-damage-classifier",
"filename": "car-damage-classifier.onnx",
},
"fusion_onnx": {
"repo_id": "junaid17/best_fusion_model_fp16",
"filename": "fusion_model.onnx",
},
"yolo_onnx": {
"repo_id": "junaid17/Yolo_Model",
"filename": "damage_detector.onnx",
},
"resnet_pt": {
"repo_id": "junaid17/car-damage-classifier",
"filename": "car-damage-classifier.pt",
},
}
def get_checkpoint_path(model_key: str) -> Path:
if model_key not in MODEL_CONFIG:
raise ValueError(f"Unknown model key: {model_key}")
config = MODEL_CONFIG[model_key]
try:
logger.info(f"Fetching {model_key} model from Hugging Face Hub...")
logger.info(f"Repo: {config['repo_id']} | File: {config['filename']}")
local_path = hf_hub_download(
repo_id=config["repo_id"],
filename=config["filename"],
)
logger.info(f"{model_key} model downloaded to: {local_path}")
return Path(local_path)
except Exception as e:
logger.exception(f"Failed to fetch {model_key} model.")
raise RuntimeError(f"Failed to load {model_key} checkpoint: {str(e)}")
class ModelLoader:
def __init__(self):
logger.info("Initializing ModelLoader...")
def get_model_path(self, model_key: str) -> Path:
return get_checkpoint_path(model_key)
def initialize_models(class_map):
logger.info("Starting model initialization pipeline...")
try:
# 1. Download/Fetch all 4 model file paths
resnet_onnx_path = get_checkpoint_path("resnet_onnx")
fusion_onnx_path = get_checkpoint_path("fusion_onnx")
yolo_onnx_path = get_checkpoint_path("yolo_onnx")
resnet_pt_path = get_checkpoint_path("resnet_pt")
# Define ONNX Execution Providers (GPU if available, fallback to CPU)
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
# 2. Initialize ONNX Runtime Sessions for Classifier Models
logger.info("Initializing ResNet ONNX Session...")
resnet_onnx_session = ort.InferenceSession(str(resnet_onnx_path), providers=providers)
logger.info("Initializing Fusion ONNX Session...")
fusion_onnx_session = ort.InferenceSession(str(fusion_onnx_path), providers=providers)
# 3. Initialize YOLO ONNX Model via Ultralytics (enables ONNX runtime with .predict API)
logger.info("Initializing YOLO ONNX Model via Ultralytics...")
yolo_onnx_model = YOLO(str(yolo_onnx_path), task="detect")
# 4. Initialize PyTorch ResNet Predictor (specifically reserved for Grad-CAM)
logger.info("Initializing PyTorch ResNet model for Grad-CAM...")
resnet_gradcam_predictor = ResnetCarDamagePredictor(
checkpoint_path=resnet_pt_path,
class_map=class_map
)
logger.info("All 4 models (3 ONNX + 1 PyTorch Grad-CAM) initialized successfully.")
return {
"resnet_onnx": resnet_onnx_session,
"fusion_onnx": fusion_onnx_session,
"yolo_onnx": yolo_onnx_model,
"resnet_pt": resnet_gradcam_predictor,
}
except Exception as e:
logger.exception("Model initialization failed.")
raise RuntimeError(f"Model initialization failed: {str(e)}") |