Update src/image_classifier_app.py

src/image_classifier_app.py

@@ -1,10 +1,97 @@
 import streamlit as st
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+from torchvision.models import resnet18
 
-st.title("Upload Test")
 
-uploaded_file = st.file_uploader("Choose a file")
+# ---------------- Constants ----------------
+CIFAR10_CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
+                   'dog', 'frog', 'horse', 'ship', 'truck']
+MODEL_PATH = "resnet18_cifar10_finetuned.pth"
 
+# ---------------- Model Loader ----------------
+@st.cache_resource
+def load_model():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    st.write("Loading ResNet18...")
+    
+    model = resnet18(pretrained=False)
+
+    st.write("Modifying model...")
+    model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
+    model.maxpool = nn.Identity()
+    in_ftrs = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Linear(in_ftrs, in_ftrs),
+        nn.ReLU(),
+        nn.Dropout(p=0.5),
+        nn.Linear(in_ftrs, 10)
+    )
+
+    st.write("Loading state_dict...")
+    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+    model.to(device)
+    model.eval()
+
+    st.write("Model ready.")
+    return model, device
+
+# ---------------- Preprocessing ----------------
+def preprocess_image(image):
+    st.write('Preparing image...')
+    transform = transforms.Compose([
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
+                             std=[0.2023, 0.1994, 0.2010])
+    ])
+    st.write('Image preparation done.')
+    return transform(image).unsqueeze(0)
+
+# ---------------- UI ----------------
+st.title("🎯 CIFAR-10 Image Classifier")
+st.write("Upload an image to classify it.")
+
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 st.write("uploaded_file value:", uploaded_file)
 
 if uploaded_file:
-    st.success("File detected")
+    try:
+        st.write('Converting image...')
+        image = Image.open(uploaded_file).convert('RGB')
+
+        st.write('Showing image...')
+        st.image(image, caption="Uploaded Image", width=200)
+
+        st.write('skipping Loading model...')
+        model, device = load_model()
+        st.write('Model loaded.')
+
+        st.write("Classifying image...")
+        with st.spinner("Classifying..."):
+            tensor = preprocess_image(image).to(device)
+            with torch.no_grad():
+                outputs = model(tensor)
+                probabilities = torch.softmax(outputs, dim=1)
+                confidence, predicted = torch.max(probabilities, 1)
+
+        st.success(f"Predicted: {CIFAR10_CLASSES[predicted.item()]}")
+        st.info(f"Confidence: {confidence.item()*100:.2f}%")
+
+    except Exception as e:
+        import traceback
+        st.error("An error occurred:")
+        st.text(traceback.format_exc())
+
+
+    top5_probs, top5_indices = torch.topk(probabilities, 5)
+    st.subheader("Top 5 Predictions")
+    for i in range(5):
+        label = CIFAR10_CLASSES[top5_indices[0][i].item()]
+        prob = top5_probs[0][i].item() * 100
+        st.write(f"{i+1}. {label} – {prob:.2f}%")
+
+    st.write("Done.")
+