Update app.py

app.py

@@ -4,4 +4,50 @@ from torchvision import transforms
 from PIL import Image
 
 
-st.write("Hi Jay and Suneel")
+st.write("Hi Jay and Suneel")
+
+# Load the image and apply transformations
+
+def predict_image(image_path, model):
+
+    image = Image.open(image_path).convert('RGB')
+    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])
+    ])
+    input_image = transform(image).unsqueeze(0)  # Add batch dimension
+
+    # Move input tensor to the device (GPU if available)
+    input_image = input_image.to(device)
+
+    # Perform inference
+    model.eval()
+    with torch.no_grad():
+        output = model(input_image)
+
+    # Get predicted class probabilities and class index
+    probabilities = torch.softmax(output, dim=1)[0]
+    predicted_class_index = torch.argmax(probabilities).item()
+
+    # Map class index to class label
+    class_labels = dataset.classes
+    predicted_class_label = class_labels[predicted_class_index]
+
+    return predicted_class_label
+#     print("Class probabilities:")
+#     for i, prob in enumerate(probabilities):
+#         print(f"{class_labels[i]}: {prob:.4f}")
+
+
+model_loaded = torchvision.models.resnet18(pretrained=False)  # Initialize ResNet18 without pretraining 
+model_loaded.fc = torch.nn.Linear(model_loaded.fc.in_features, len(dataset.classes))  # Modify the fully connected layer
+model_loaded = model_loaded.to(device)  # Move the model to the appropriate device (GPU or CPU)
+
+# Load the saved state dictionary into the model
+model_path = 'resnet18_custom_model.pth'
+model_loaded.load_state_dict(torch.load(model_path))
+
+# Set the model to evaluation mode
+model_loaded.eval()