dd

app.py

@@ -3,12 +3,30 @@ import gradio as gr
 
 import torch
 import torchvision
+import timm
 
 checkpoint = torch.load('v4-epoch=19-val_loss=0.6964-val_accuracy=0.8964.ckpt', map_location=torch.device('cpu'))
 state_dict = checkpoint["state_dict"]                                                                                                                                                                                                         
 model_weights = state_dict                                                                                                                                                                                                                    
 for key in list(model_weights):                                                                                                                                                                                                               
     model_weights[key.replace("backbone.", "")] = model_weights.pop(key)                                                                                                                                                                      
+
+
+def get_model():
+    model = timm.create_model('tf_efficientnet_b1', pretrained=True, num_classes=2, global_pool='catavgmax')
+    num_in_features = model.get_classifier().in_features
+    from torch import nn
+
+    model.fc = nn.Sequential(
+        nn.Linear(in_features=num_in_features, out_features=1024, bias=False),
+        nn.ReLU(),
+        nn.Linear(in_features=1024, out_features=2, bias=False),
+    )
+
+    return model
+
+model = get_model()
+
 model.load_state_dict(model_weights).eval()
 
 import requests