Spaces:
Runtime error
Runtime error
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| # setup Device to CUDA | |
| device='cuda' if torch.cuda.is_available() else 'cpu' | |
| device | |
| class CustomBlock(nn.Module): | |
| def __init__(self, in_channels, out_channels, stride=1): | |
| super(CustomBlock, self).__init__() | |
| self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1) | |
| self.bn1 = nn.BatchNorm2d(out_channels) | |
| self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1) | |
| self.bn2 = nn.BatchNorm2d(out_channels) | |
| self.conv3 = nn.Conv2d(out_channels, out_channels , kernel_size=1) | |
| self.bn3 = nn.BatchNorm2d(out_channels) | |
| self.relu = nn.ReLU() | |
| self.identity_conv=nn.Conv2d(in_channels,out_channels,kernel_size=1, stride=stride, padding=1) | |
| def forward(self, x): | |
| identity = x | |
| #print(identity.shape) | |
| x = self.conv1(x) | |
| x = self.bn1(x) | |
| x = self.relu(x) | |
| x = self.conv2(x) | |
| x = self.bn2(x) | |
| x = self.relu(x) | |
| x = self.conv3(x) | |
| x = self.bn3(x) | |
| #print(x.shape) | |
| if self.identity_conv is not None: | |
| identity = self.identity_conv(identity) | |
| if x.shape != identity.shape: | |
| identity=nn.functional.interpolate(identity,size=(x.shape[2],x.shape[3]),mode='nearest') | |
| x += identity | |
| x = self.relu(x) | |
| return x | |
| class SimpleResNet(nn.Module): | |
| def __init__(self, num_classes=13): | |
| super(SimpleResNet, self).__init__() | |
| self.conv1 = nn.Conv2d(3,16, kernel_size=3, stride=1, padding=1) | |
| self.bn1 = nn.BatchNorm2d(16) | |
| self.relu = nn.ReLU() | |
| self.maxpool = nn.MaxPool2d(kernel_size=2, stride=1, padding=1) | |
| self.block1 = CustomBlock(16, 32) | |
| self.block2 = CustomBlock(32,64) | |
| self.block3 = CustomBlock(64,128) | |
| self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) | |
| self.flatten=nn.Flatten() | |
| self.fc = nn.Linear(128 ,128) | |
| self.fc2=nn.Linear(128,256) | |
| self.drop=nn.Dropout(p=0.5) | |
| self.fc3=nn.Linear(256,num_classes) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = self.bn1(x) | |
| x = self.relu(x) | |
| x = self.maxpool(x) | |
| x = self.block1(x) | |
| x = self.block2(x) | |
| x = self.block3(x) | |
| x = self.avgpool(x) | |
| x = self.flatten(x) | |
| x = self.fc(x) | |
| x=self.fc2(x) | |
| x=self.drop(x) | |
| x=self.fc3(x) | |
| return x | |
| model=SimpleResNet(num_classes=21).to(device) | |
| model.load_state_dict(torch.load('model_with_info_path_final.pt',map_location=device)) | |
| model.eval() | |
| classes=['206','207','405','Dena','L90','Mazda-vanet','Naisan','Pars','Paykan-Vanet','Pride','Pride_vanet','Quiek', | |
| 'Saina','Tiba','Truck-Benz','Truck-Renault','Unknown','Volvo-FH-FM','Volvo-N10','Volvo-NH','samand'] | |
| transform=transforms.Compose([transforms.Resize((224,224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((.5),(.5))]) | |
| def classify_image(img1): | |
| model.eval() | |
| with torch.inference_mode(): | |
| #img1=Image.open(img1).convert("RGB") | |
| img1=transform(img1).unsqueeze(0).to(device) | |
| y_logits=model(img1) | |
| y_pred=torch.softmax(y_logits,dim=1)#.argmax(dim=1) | |
| conf,pred_class=torch.max(y_pred,dim=1) | |
| if conf.item()<0.55: | |
| return f"I'm not sure what this is and confidence:{conf.item():.2f}" | |
| else: | |
| return f'Car: {classes[pred_class]} confidence:{conf.item():.2f}' | |
| #img1=Image.open(img1).convert("RGB") | |
| #img1=transform(img1).unsqueeze(0).to(device) | |
| #print(img.shape) | |
| # y_logits=model(img1) | |
| #y_pred=torch.softmax(y_logits,dim=1).argmax(dim=1) | |
| #confidence = {classes[i]: float(y_pred[i]) for i in range(len(classes))} | |
| #return classes[pred_class] | |
| interface = gr.Interface( | |
| fn=classify_image, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Label(num_top_classes=21), | |
| title="Iranian Car Classifier") | |
| interface.launch(share=True) |