initial commit

.gitignore

@@ -0,0 +1,3 @@
+venv/
+venv
+__pycache__

app.py

@@ -0,0 +1,32 @@
+import torch
+import torchvision
+from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
+
+from torch import nn
+
+from PIL import Image
+from model import create_effnet_v2_model
+
+class_names = ['Honda', 'Hyundai', 'Toyota']
+
+effnet_v2, transforms = create_effnet_v2_model(num_classes=len(class_names), weights_path="efficient_net_s_carvision_3.pth")
+
+def predict(model, image_path, device):
+
+    image = Image.open(image_path)
+    image = transforms(image).unsqueeze(0)
+    image = image.to(device)
+    output = model(image)
+    
+    model.eval()
+    with torch.inference_mode():
+        probs = torch.softmax(output, dim=1)
+
+    pred_labels_and_probs = {class_names[i]: float(probs[0, i]) for i in range(len(class_names))}
+
+
+    return pred_labels_and_probs
+
+print(predict(effnet_v2, "examples/Toyota_Tacoma_2017_36_18_270_35_6_75_70_212_19_RWD_5_4_Pickup_xQa.jpg", torch.device("cpu")))
+
+# print(predict(effnet_v2, "test.jpg", torch.device("cuda:0")))

efficient_net_s_carvision_3.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ca71eab330a5ea4163fea0751b4c6551bfaf82461ebb751e6128cf3afb4b738
+size 81639653

model.py

@@ -0,0 +1,19 @@
+import torch
+import torchvision
+
+from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
+from torch import nn
+
+
+def create_effnet_v2_model(weights_path, num_classes=3):
+    weights = EfficientNet_V2_S_Weights.DEFAULT
+    transforms = weights.transforms()
+
+    model = efficientnet_v2_s()
+    model.classifier = nn.Sequential(
+        nn.Dropout(0.0),
+        nn.Linear(in_features=1280, out_features=num_classes)
+    )
+    model.load_state_dict(torch.load(f = weights_path, map_location=torch.device("cpu")))
+
+    return model, transforms