Sher1988 commited on
Commit
5588fb5
·
verified ·
1 Parent(s): 6dcb912

Delete src

Browse files
Files changed (2) hide show
  1. src/image_classifier_app.py +0 -93
  2. src/streamlit_app.py +0 -40
src/image_classifier_app.py DELETED
@@ -1,93 +0,0 @@
1
- import streamlit as st
2
- import torch
3
- import torch.nn as nn
4
- import torchvision.transforms as transforms
5
- from PIL import Image
6
- from torchvision.models import resnet18
7
- import os
8
-
9
- # Get the directory where the current script (app.py) is located
10
- # Since app.py is in /app/src/ and the model is in /app/
11
- BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
12
- MODEL_PATH = os.path.join(BASE_DIR, "resnet18_cifar10_finetuned.pth")
13
-
14
- # Use MODEL_PATH in your load_model function
15
- # Example: model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
16
-
17
- # ---------------- Constants ----------------
18
- CIFAR10_CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
19
-
20
- # ---------------- Model Loader ----------------
21
- @st.cache_resource
22
- def load_model():
23
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
24
-
25
- model = resnet18(pretrained=False)
26
-
27
- model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
28
- model.maxpool = nn.Identity()
29
- in_ftrs = model.fc.in_features
30
- model.fc = nn.Sequential(
31
- nn.Linear(in_ftrs, in_ftrs),
32
- nn.ReLU(),
33
- nn.Dropout(p=0.5),
34
- nn.Linear(in_ftrs, 10)
35
- )
36
-
37
- model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
38
- model.to(device)
39
- model.eval()
40
-
41
- return model, device
42
-
43
- # ---------------- Preprocessing ----------------
44
- def preprocess_image(image):
45
- transform = transforms.Compose([
46
- transforms.Resize((32, 32)),
47
- transforms.ToTensor(),
48
- transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
49
- std=[0.2023, 0.1994, 0.2010])
50
- ])
51
- return transform(image).unsqueeze(0)
52
-
53
- # ---------------- UI ----------------
54
- st.title("🎯 CIFAR-10 Image Classifier")
55
- st.write("Upload an image to classify it.")
56
- st.write("ResNet18 model finetuned for 3 epochs with 95.6% accuracy on CIFAR10 images.")
57
- st.write("The model performs well on images from CIFAR10 dataset.")
58
-
59
- uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
60
-
61
- if uploaded_file:
62
- try:
63
- image = Image.open(uploaded_file).convert('RGB')
64
-
65
- st.image(image, caption="Uploaded Image", width=200)
66
-
67
- model, device = load_model()
68
-
69
- with st.spinner("Classifying..."):
70
- tensor = preprocess_image(image).to(device)
71
- with torch.no_grad():
72
- outputs = model(tensor)
73
- probabilities = torch.softmax(outputs, dim=1)
74
- confidence, predicted = torch.max(probabilities, 1)
75
-
76
- st.success(f"Predicted: {CIFAR10_CLASSES[predicted.item()]}")
77
- st.info(f"Confidence: {confidence.item()*100:.2f}%")
78
-
79
- except Exception as e:
80
- import traceback
81
- st.error("An error occurred:")
82
- st.text(traceback.format_exc())
83
-
84
-
85
- top5_probs, top5_indices = torch.topk(probabilities, 5)
86
- st.subheader("Top 5 Predictions")
87
- for i in range(5):
88
- label = CIFAR10_CLASSES[top5_indices[0][i].item()]
89
- prob = top5_probs[0][i].item() * 100
90
- st.write(f"{i+1}. {label} – {prob:.2f}%")
91
-
92
- st.write("Done.")
93
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/streamlit_app.py DELETED
@@ -1,40 +0,0 @@
1
- import altair as alt
2
- import numpy as np
3
- import pandas as pd
4
- import streamlit as st
5
-
6
- """
7
- # Welcome to Streamlit!
8
-
9
- Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
10
- If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
11
- forums](https://discuss.streamlit.io).
12
-
13
- In the meantime, below is an example of what you can do with just a few lines of code:
14
- """
15
-
16
- num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
17
- num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
18
-
19
- indices = np.linspace(0, 1, num_points)
20
- theta = 2 * np.pi * num_turns * indices
21
- radius = indices
22
-
23
- x = radius * np.cos(theta)
24
- y = radius * np.sin(theta)
25
-
26
- df = pd.DataFrame({
27
- "x": x,
28
- "y": y,
29
- "idx": indices,
30
- "rand": np.random.randn(num_points),
31
- })
32
-
33
- st.altair_chart(alt.Chart(df, height=700, width=700)
34
- .mark_point(filled=True)
35
- .encode(
36
- x=alt.X("x", axis=None),
37
- y=alt.Y("y", axis=None),
38
- color=alt.Color("idx", legend=None, scale=alt.Scale()),
39
- size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
40
- ))