Spaces:
Paused
Paused
| import logging | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from torchvision import transforms, models | |
| from PIL import Image, UnidentifiedImageError | |
| from transformers import ConvNextImageProcessor | |
| import onnxruntime as ort | |
| # ---------------- LOGGING SETUP ---------------- | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s - %(levelname)s - %(name)s - %(message)s" | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # ===================================================================== | |
| # 1. PYTORCH RESNET MODEL & PREDICTOR (RETAINED EXCLUSIVELY FOR GRAD-CAM) | |
| # ===================================================================== | |
| class Car_Classifier_Resnet(nn.Module): | |
| def __init__(self, num_classes): | |
| super().__init__() | |
| logger.info("Initializing ResNet18 PyTorch architecture for Grad-CAM...") | |
| self.model = models.resnet18(weights="DEFAULT") | |
| for param in self.model.parameters(): | |
| param.requires_grad = False | |
| for param in self.model.layer3.parameters(): | |
| param.requires_grad = True | |
| for param in self.model.layer4.parameters(): | |
| param.requires_grad = True | |
| self.model.fc = nn.Sequential( | |
| nn.Dropout(0.5), | |
| nn.Linear(self.model.fc.in_features, 256), | |
| nn.ReLU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(256, num_classes) | |
| ) | |
| def forward(self, x): | |
| return self.model(x) | |
| class ResnetCarDamagePredictor: | |
| """PyTorch Predictor wrapper used strictly by Grad-CAM for backpropagation.""" | |
| def __init__(self, checkpoint_path, class_map): | |
| logger.info("Initializing PyTorch ResNet predictor for Grad-CAM...") | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.class_map = class_map | |
| self.test_transforms = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| [0.485, 0.456, 0.406], | |
| [0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| try: | |
| self.model = Car_Classifier_Resnet(num_classes=len(class_map)) | |
| checkpoint = torch.load(checkpoint_path, map_location=self.device) | |
| state_dict = checkpoint.get("model_state_dict", checkpoint) | |
| self.model.load_state_dict(state_dict) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| logger.info("PyTorch ResNet model loaded successfully for Grad-CAM.") | |
| except Exception as e: | |
| logger.exception("Failed to load PyTorch ResNet model.") | |
| raise RuntimeError(f"Failed to load ResNet PyTorch model: {str(e)}") | |
| # ===================================================================== | |
| # 2. ONNX RESNET PREDICTOR (FAST INFERENCE) | |
| # ===================================================================== | |
| class ResnetONNXPredictor: | |
| def __init__(self, session: ort.InferenceSession, class_map: dict): | |
| logger.info("Initializing ResNet ONNX Predictor...") | |
| self.session = session | |
| self.class_map = class_map | |
| self.input_name = self.session.get_inputs()[0].name | |
| self.transforms = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| [0.485, 0.456, 0.406], | |
| [0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| def predict(self, image_input): | |
| logger.info("Starting ResNet ONNX prediction...") | |
| try: | |
| if isinstance(image_input, str): | |
| image = Image.open(image_input).convert("RGB") | |
| elif isinstance(image_input, Image.Image): | |
| image = image_input.convert("RGB") | |
| else: | |
| raise TypeError("image_input must be a file path or PIL.Image") | |
| # Preprocess to [1, 3, 224, 224] NumPy array | |
| input_tensor = self.transforms(image).unsqueeze(0).numpy() | |
| # Execute ONNX forward pass | |
| outputs = self.session.run(None, {self.input_name: input_tensor}) | |
| logits = outputs[0][0] | |
| # Softmax calculation | |
| exp_logits = np.exp(logits - np.max(logits)) | |
| probs = exp_logits / np.sum(exp_logits) | |
| class_probs = { | |
| self.class_map[i]: float(probs[i]) | |
| for i in range(len(self.class_map)) | |
| } | |
| sorted_probs = dict( | |
| sorted(class_probs.items(), key=lambda x: x[1], reverse=True) | |
| ) | |
| return sorted_probs | |
| except UnidentifiedImageError: | |
| logger.error("Invalid image file provided to ResNet ONNX predictor.") | |
| raise ValueError("Invalid image file provided") | |
| except Exception as e: | |
| logger.exception("ResNet ONNX prediction failed.") | |
| raise RuntimeError(f"ResNet ONNX prediction failed: {str(e)}") | |
| # ===================================================================== | |
| # 3. ONNX FUSION PREDICTOR (EFFICIENTNET + CONVNEXT) | |
| # ===================================================================== | |
| class FusionONNXPredictor: | |
| def __init__(self, session: ort.InferenceSession, class_map: dict, convnext_model_name="facebook/convnext-small-224"): | |
| logger.info("Initializing Fusion ONNX Predictor...") | |
| self.session = session | |
| self.class_map = class_map | |
| # Dynamically retrieve ONNX input names (EffNet & ConvNeXt) | |
| self.input_eff_name = self.session.get_inputs()[0].name | |
| self.input_cnx_name = self.session.get_inputs()[1].name | |
| self.eff_normalize = transforms.Compose([ | |
| transforms.Resize((260, 260)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| [0.485, 0.456, 0.406], | |
| [0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| self.convnext_processor = ConvNextImageProcessor.from_pretrained(convnext_model_name) | |
| def predict(self, image_input): | |
| logger.info("Starting Fusion ONNX prediction...") | |
| try: | |
| if isinstance(image_input, str): | |
| image = Image.open(image_input).convert("RGB") | |
| elif isinstance(image_input, Image.Image): | |
| image = image_input.convert("RGB") | |
| else: | |
| raise TypeError("image_input must be a file path or PIL.Image") | |
| # EfficientNet preprocessing -> [1, 3, 260, 260] | |
| pixel_eff_np = self.eff_normalize(image).unsqueeze(0).numpy() | |
| # ConvNeXt preprocessing -> [1, 3, 224, 224] | |
| inputs_cnx = self.convnext_processor(images=image, return_tensors="pt") | |
| pixel_cnx_np = inputs_cnx["pixel_values"].numpy() | |
| # Execute dual-input ONNX forward pass | |
| onnx_inputs = { | |
| self.input_eff_name: pixel_eff_np, | |
| self.input_cnx_name: pixel_cnx_np | |
| } | |
| outputs = self.session.run(None, onnx_inputs) | |
| logits = outputs[0][0] | |
| # Softmax calculation | |
| exp_logits = np.exp(logits - np.max(logits)) | |
| probs = exp_logits / np.sum(exp_logits) | |
| class_probs = { | |
| self.class_map[i]: float(probs[i]) | |
| for i in range(len(self.class_map)) | |
| } | |
| sorted_probs = dict( | |
| sorted(class_probs.items(), key=lambda x: x[1], reverse=True) | |
| ) | |
| return sorted_probs | |
| except UnidentifiedImageError: | |
| logger.error("Invalid image file provided to Fusion ONNX predictor.") | |
| raise ValueError("Invalid image file provided") | |
| except Exception as e: | |
| logger.exception("Fusion ONNX prediction failed.") | |
| raise RuntimeError(f"Fusion ONNX prediction failed: {str(e)}") |