Create utils.py

utils.py

@@ -0,0 +1,27 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+
+def load_image(image_path):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                           std=[0.229, 0.224, 0.225])
+    ])
+    
+    image = Image.open(image_path).convert('RGB')
+    return transform(image).unsqueeze(0)
+
+def predict_toxicity(model, image_tensor, device):
+    model.eval()
+    with torch.no_grad():
+        image_tensor = image_tensor.to(device)
+        outputs = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        prediction = torch.argmax(probabilities, dim=1)
+    return prediction.item(), probabilities[0].cpu().numpy()
+
+def get_label(prediction):
+    return "Toxic" if prediction == 1 else "Non-Toxic"