Spaces:
Sleeping
Sleeping
File size: 7,172 Bytes
1782395 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | import os
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image, UnidentifiedImageError
from transformers import DeiTForImageClassification, DeiTImageProcessor
# ================================ ResNet-18 Classifier ================================
class Car_Classifier_Resnet(nn.Module):
def __init__(self, num_classes):
super().__init__()
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
# Replace FC head
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:
def __init__(self, checkpoint_path, class_map):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.class_map = class_map
self.test_transforms = transforms.Compose([
transforms.Resize((128, 128)),
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)
self.model.load_state_dict(checkpoint["model_state_dict"])
self.model.to(self.device)
self.model.eval()
except Exception as e:
raise RuntimeError(f"Failed to load ResNet model: {str(e)}")
def resnet_predict(self, image_input):
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")
image = self.test_transforms(image)
image = image.unsqueeze(0).to(self.device)
with torch.no_grad():
outputs = self.model(image)
probs = torch.nn.functional.softmax(outputs, dim=1)[0]
class_probs = {
self.class_map[i]: float(probs[i].item())
for i in range(len(self.class_map))
}
return dict(sorted(class_probs.items(), key=lambda x: x[1], reverse=True))
except UnidentifiedImageError:
raise ValueError("Invalid image file provided")
except Exception as e:
raise RuntimeError(f"ResNet prediction failed: {str(e)}")
# ================================ DeiT Classifier ================================
class DeitCarDamagePredictor:
def __init__(self, checkpoint_path, class_map):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.class_map = class_map
self.checkpoint_path = checkpoint_path
self.transform = transforms.Compose([transforms.Resize((224, 224))])
model_name = "facebook/deit-base-distilled-patch16-224"
try:
self.processor = DeiTImageProcessor.from_pretrained(model_name)
self.model = DeiTForImageClassification.from_pretrained(
model_name,
num_labels=len(class_map),
ignore_mismatched_sizes=True
)
checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
self.model.load_state_dict(checkpoint["model_state_dict"])
self.model.to(self.device)
self.model.eval()
except Exception as e:
raise RuntimeError(f"Failed to load DeiT model: {str(e)}")
def deit_predict(self, image_input):
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")
image = self.transform(image)
inputs = self.processor(image, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
class_probs = {
self.class_map[i]: float(probs[i].item())
for i in range(len(self.class_map))
}
return dict(sorted(class_probs.items(), key=lambda x: x[1], reverse=True))
except UnidentifiedImageError:
raise ValueError("Invalid image file provided")
except Exception as e:
raise RuntimeError(f"DeiT prediction failed: {str(e)}")
# ================================ Fusion Predictor ================================
class FusionCarDamagePredictor:
def __init__(self, resnet_predictor, deit_predictor, resnet_weight=0.5, deit_weight=0.5):
if resnet_weight < 0 or deit_weight < 0:
raise ValueError("Weights must be non-negative")
total = resnet_weight + deit_weight
if total == 0:
raise ValueError("At least one weight must be greater than 0")
self.resnet_predictor = resnet_predictor
self.deit_predictor = deit_predictor
self.resnet_weight = resnet_weight / total
self.deit_weight = deit_weight / total
def fuse_predict(self, image_input):
try:
resnet_output = self.resnet_predictor.resnet_predict(image_input)
deit_output = self.deit_predictor.deit_predict(image_input)
all_classes = set(resnet_output.keys()).union(set(deit_output.keys()))
fused_output = {}
for cls in all_classes:
resnet_prob = resnet_output.get(cls, 0.0)
deit_prob = deit_output.get(cls, 0.0)
fused_prob = self.resnet_weight * resnet_prob + self.deit_weight * deit_prob
fused_output[cls] = float(fused_prob)
fused_output = dict(sorted(fused_output.items(), key=lambda x: x[1], reverse=True))
final_class = next(iter(fused_output))
final_confidence = fused_output[final_class]
return {
"resnet_output": resnet_output,
"deit_output": deit_output,
"fused_output": fused_output,
"final_prediction": final_class,
"final_confidence": round(final_confidence, 4)
}
except Exception as e:
raise RuntimeError(f"Fusion prediction failed: {str(e)}") |