File size: 1,332 Bytes
138a538 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
# ✅ Load Class Names
with open("class_names.txt", "r") as f:
class_names = [line.strip() for line in f.readlines()]
# ✅ Load Trained Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.mobilenet_v2(pretrained=False)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(class_names))
model.load_state_dict(torch.load("plant_disease_model.pth", map_location=device))
model = model.to(device)
model.eval()
# ✅ Image Transformations (Must match training settings)
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
# ✅ Function to Make Predictions
def predict_image(image_path):
image = Image.open(image_path).convert("RGB")
image = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
output = model(image)
predicted_class = torch.argmax(output, dim=1).item()
return class_names[predicted_class]
# ✅ Test the model (optional)
if __name__ == "__main__":
sample_image = "test_image.jpg" # Replace with an actual image path
prediction = predict_image(sample_image)
print(f"Predicted Class: {prediction}")
|