| from PIL import Image |
| import torch |
| from torchvision import transforms |
| from model import load_model |
|
|
| |
| model1 = load_model("m1.safetensors") |
| model2 = load_model("m2.safetensors") |
| model3 = load_model("m3.safetensors") |
|
|
| models = [model1, model2, model3] |
|
|
| transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.ToTensor() |
| ]) |
|
|
| def predict(image): |
| img = Image.open(image).convert("RGB") |
| x = transform(img).unsqueeze(0) |
|
|
| probs = [] |
|
|
| with torch.no_grad(): |
| for model in models: |
| out = model(x) |
| prob = torch.sigmoid(out).item() |
| probs.append(prob) |
|
|
| |
| final_prob = sum(probs) / len(probs) |
|
|
| return final_prob |
