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)