Update app.py

app.py

@@ -1,46 +1,46 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-from torchvision import transforms, datasets, models
-from PIL import Image
-
-# Title
-st.title("Brain Tumor Classification")
-
-# Class names
-class_names = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']
-
-# Load pre-trained ResNet18 model
-model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
-num_of_classes = len(class_names)
-num_of_features = model.fc.in_features
-model.fc = nn.Linear(num_of_features, num_of_classes)
-
-# Load trained model weights
-model.load_state_dict(torch.load('resnet18_model (1).pth', map_location=torch.device('cpu')))
-model.eval()
-
-# Image upload
-uploaded_img = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-if uploaded_img is not None:
-    # Display uploaded image
-    image = Image.open(uploaded_img)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    # Image transformations
-    sample_transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.1776, 0.1776, 0.1776], std=[0.1735, 0.1735, 0.1735])
-    ])
-
-    # Apply transformations
-    transformed_img = sample_transform(image).unsqueeze(0)
-
-    # Model inference
-    with torch.no_grad():
-        pred = model(transformed_img).argmax(dim=1).item()
-
-    # Display prediction
-    st.success(f"Predicted Class: {class_names[pred]}")
+import streamlit as st
+import torch
+import torch.nn as nn
+from torchvision import transforms, datasets, models
+from PIL import Image
+
+# Title
+st.title("Brain Tumor Classification")
+
+# Class names
+class_names = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']
+
+# Load pre-trained ResNet18 model
+model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+num_of_classes = len(class_names)
+num_of_features = model.fc.in_features
+model.fc = nn.Linear(num_of_features, num_of_classes)
+
+# Load trained model weights
+model.load_state_dict(torch.load('resnet18_model (1).pth', map_location=torch.device('cpu')))
+model.eval()
+
+# Image upload
+uploaded_img = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+if uploaded_img is not None:
+    # Display uploaded image
+    image = Image.open(uploaded_img)
+    st.image(image, caption="Uploaded Image", use_container_width =True)
+
+    # Image transformations
+    sample_transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.1776, 0.1776, 0.1776], std=[0.1735, 0.1735, 0.1735])
+    ])
+
+    # Apply transformations
+    transformed_img = sample_transform(image).unsqueeze(0)
+
+    # Model inference
+    with torch.no_grad():
+        pred = model(transformed_img).argmax(dim=1).item()
+
+    # Display prediction
+    st.success(f"Predicted Class: {class_names[pred]}")