Upload inference.py

inference.py

@@ -0,0 +1,48 @@
+import json
+import time
+from PIL import Image
+import torch
+from torchvision.transforms import transforms
+
+model = torch.load('/path/to/your/model.pth').to("cuda")
+model.eval()
+transform = transforms.Compose([
+    transforms.Resize((448, 448)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[
+        0.48145466,
+        0.4578275,
+        0.40821073
+    ], std=[
+        0.26862954,
+        0.26130258,
+        0.27577711
+    ])
+])
+
+with open("tags_8041.json", "r") as file:
+    tags = json.load(file)
+allowed_tags = sorted(tags)
+allowed_tags.insert(0, "placeholder0")
+allowed_tags.append("placeholder1")
+allowed_tags.append("explicit")
+allowed_tags.append("questionable")
+allowed_tags.append("safe")
+
+image_path = "/path/to/your/image.jpg"
+start = time.time()
+img = Image.open(image_path).convert('RGB')
+tensor = transform(img).unsqueeze(0).to("cuda")  # transform and add batch dimension
+
+with torch.no_grad():
+    out = model(tensor)
+probabilities = torch.nn.functional.sigmoid(out[0])
+indices = torch.where(probabilities > 0.3)[0]
+values = probabilities[indices]
+
+for i in range(indices.size(0)):
+    print(allowed_tags[indices[i]], values[i].item())
+
+end = time.time()
+print(f'Executed in {end - start} seconds')
+print("\n\n", end="")