Upload 2 files

7704_inference.py

@@ -9,9 +9,17 @@ model = torch.load('model.pth').to("cuda")
 model.eval()
 config = resolve_data_config({}, model=model)
 transform = transforms.Compose([
-                transforms.Resize((384, 384)),
+                transforms.Resize((224, 224)),
                 transforms.ToTensor(),
-                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+                transforms.Normalize(mean=[
+        0.48145466,
+        0.4578275,
+        0.40821073
+    ], std=[
+        0.26862954,
+        0.26130258,
+        0.27577711
+    ])  # Normalize image
             ])
 
 with open("tags.json", "r") as file:
@@ -35,4 +43,4 @@ for i in range(top10_prob.size(0)):
     print(allowed_tags[top10_catid[i]], top10_prob[i].item())
 end = time.time()
 print(f'Executed in {end - start} seconds')
-print("\n\n", end="")
+print("\n\n", end="")