torch model

app.py

@@ -1,18 +1,45 @@
 import gradio as gr
-from fastai.vision.all import *
+import torch
+from PIL import Image
+from torchvision import transforms
+import torchvision.models as models
+import torch.nn as nn
 
-learn = load_learner('model.pkl')
-categories = ['911', 'Cayenne', 'Cayman', 'Macan', 'Panamera', 'Taycan']
+# Define the class names
+class_names = ['911', 'cayenne', 'cayman', 'macan', 'panamera', 'taycan']
 
-def classify_image(img):
-    img = PILImage.create(img) 
-    dl = learn.dls.test_dl([img])
-    probs, _ = learn.tta(dl=dl)
-    return dict(zip(categories, map(float, probs[0])))
+# Instantiate the model and load state_dict
+model_ft = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
+for param in model_ft.parameters():
+    param.requires_grad = False
+for param in model_ft.layer4.parameters():
+    param.requires_grad = True
 
+num_ftrs = model_ft.fc.in_features
+model_ft.fc = nn.Linear(num_ftrs, len(class_names))
+model_ft = model_ft.to('cuda' if torch.cuda.is_available() else 'cpu')
+model_ft.load_state_dict(torch.load('model_ft.pth'))
+model_ft.eval()
 
-image = gr.Image(width=224, height=224)
-label = gr.Label()
+# Define preprocessing transforms
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
 
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label)
-intf.launch(inline=False, share=True)
+# Define the prediction function
+def predict(image):
+    image = preprocess(image).unsqueeze(0).to(model_ft.device)  # Add batch dimension and move to device
+    with torch.no_grad():
+        outputs = model_ft(image)
+        _, predicted = torch.max(outputs, 1)
+    return class_names[predicted.item()]
+
+# Create Gradio interface
+iface = gr.Interface(fn=predict,
+                     inputs=gr.inputs.Image(type="pil"),
+                     outputs="text")
+
+iface.launch()

model.pkl → porsche_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6918bdfe8f5c78dd1e0427173cb35f60e26d9aff6ebcb3bd502fcaaa2ada966b
-size 87485098
+oid sha256:9506b43b2d24c94d97a2e4767cb1b05f84132373579fc7bdd8fb775bbc9dc9c5
+size 85292402