Spaces:
Sleeping
Sleeping
Upload 5 files
Browse files- app.py +90 -0
- class_names.txt +9 -0
- cnn.pth +3 -0
- model.py +91 -0
- requirements.txt +3 -0
app.py
ADDED
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
### 1. Imports and class names setup ###
|
2 |
+
import gradio as gr
|
3 |
+
import os
|
4 |
+
import torch
|
5 |
+
from model import TinyCNN
|
6 |
+
from timeit import default_timer as timer
|
7 |
+
from typing import Tuple, Dict
|
8 |
+
import torch
|
9 |
+
import torchvision
|
10 |
+
from torchvision import transforms
|
11 |
+
from torch import nn
|
12 |
+
# Setup class names
|
13 |
+
with open("class_names.txt", "r") as f: # reading them in from class_names.txt
|
14 |
+
class_names = [defects.strip() for defects in f.readlines()]
|
15 |
+
|
16 |
+
### 2. Model and transforms preparation ###
|
17 |
+
|
18 |
+
# Create model
|
19 |
+
TinyCNN_model = TinyCNN(input_shape=3,
|
20 |
+
hidden_units=64,
|
21 |
+
output_shape=len(class_names))
|
22 |
+
|
23 |
+
|
24 |
+
transform = transforms.Compose([
|
25 |
+
transforms.Resize((224, 224)),
|
26 |
+
transforms.ToTensor()
|
27 |
+
])
|
28 |
+
|
29 |
+
# Load saved weights
|
30 |
+
TinyCNN_model.load_state_dict(
|
31 |
+
torch.load(
|
32 |
+
f="cnn.pth",
|
33 |
+
map_location=torch.device("cpu"), # load to CPU
|
34 |
+
)
|
35 |
+
)
|
36 |
+
|
37 |
+
### 3. Predict function ###
|
38 |
+
|
39 |
+
# Create predict function
|
40 |
+
def predict(img) -> Tuple[Dict, float]:
|
41 |
+
"""Transforms and performs a prediction on img and returns prediction and time taken.
|
42 |
+
"""
|
43 |
+
# Start the timer
|
44 |
+
start_time = timer()
|
45 |
+
|
46 |
+
# Transform the target image and add a batch dimension
|
47 |
+
img = transform(img).unsqueeze(dim=0)
|
48 |
+
|
49 |
+
# Put model into evaluation mode and turn on inference mode
|
50 |
+
TinyCNN_model.eval()
|
51 |
+
with torch.inference_mode():
|
52 |
+
# Pass the transformed image through the model and turn the prediction logits into prediction probabilities
|
53 |
+
pred_probs = torch.softmax(TinyCNN_model(img), dim=1)
|
54 |
+
|
55 |
+
# Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
|
56 |
+
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
|
57 |
+
|
58 |
+
# Calculate the prediction time
|
59 |
+
pred_time = round(timer() - start_time, 5)
|
60 |
+
|
61 |
+
# Return the prediction dictionary and prediction time
|
62 |
+
return pred_labels_and_probs, pred_time
|
63 |
+
|
64 |
+
|
65 |
+
### 4. Gradio app ###
|
66 |
+
|
67 |
+
# Create title, description and article strings
|
68 |
+
title = "Expression Detection"
|
69 |
+
description = "An app to predict emotions from the list.[Angry, Disgust, Fear, Happy, Neutral, Sad, Surprise]"
|
70 |
+
article = "Created as a college project."
|
71 |
+
|
72 |
+
# Create examples list from "examples/" directory
|
73 |
+
example_list = [["examples/" + example] for example in os.listdir("examples")]
|
74 |
+
|
75 |
+
# Create Gradio interface
|
76 |
+
demo = gr.Interface(
|
77 |
+
fn=predict,
|
78 |
+
inputs=gr.Image(type="pil"),
|
79 |
+
outputs=[
|
80 |
+
gr.Label(num_top_classes=5, label="Predictions"),
|
81 |
+
gr.Image(label="Defects"),
|
82 |
+
],
|
83 |
+
examples=example_list,
|
84 |
+
title=title,
|
85 |
+
description=description,
|
86 |
+
article=article,
|
87 |
+
)
|
88 |
+
|
89 |
+
# Launch the app!
|
90 |
+
demo.launch()
|
class_names.txt
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Center
|
2 |
+
Donut
|
3 |
+
Edge-Loc
|
4 |
+
Edge-Ring
|
5 |
+
Loc
|
6 |
+
Near-full
|
7 |
+
Random
|
8 |
+
Scratch
|
9 |
+
none
|
cnn.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:be41344ebccd7a18bdaa26887f89c62324ee8b6a3032744eb0023a6b215be216
|
3 |
+
size 115453650
|
model.py
ADDED
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torchvision
|
3 |
+
|
4 |
+
from torch import nn
|
5 |
+
|
6 |
+
class TinyCNN(nn.Module):
|
7 |
+
|
8 |
+
def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
|
9 |
+
super().__init__()
|
10 |
+
self.conv_block_1 = nn.Sequential(
|
11 |
+
nn.Conv2d(in_channels=input_shape,
|
12 |
+
out_channels=hidden_units,
|
13 |
+
kernel_size=3, # how big is the square that's going over the image?
|
14 |
+
stride=1, # default
|
15 |
+
padding=1), # options = "valid" (no padding) or "same" (output has same shape as input) or int for specific number
|
16 |
+
nn.BatchNorm2d(hidden_units),
|
17 |
+
nn.ReLU(),
|
18 |
+
# nn.Conv2d(in_channels=hidden_units,
|
19 |
+
# out_channels=128,
|
20 |
+
# kernel_size=3,
|
21 |
+
# stride=1,
|
22 |
+
# padding=0),
|
23 |
+
# nn.BatchNorm2d(128),
|
24 |
+
# nn.ReLU(),
|
25 |
+
nn.MaxPool2d(kernel_size=2,
|
26 |
+
stride=2), # default stride value is same as kernel_size
|
27 |
+
nn.Dropout(p=0.25)
|
28 |
+
)
|
29 |
+
self.conv_block_2 = nn.Sequential(
|
30 |
+
nn.Conv2d(hidden_units, 128, kernel_size=3, padding=1),
|
31 |
+
# nn.ReLU(),
|
32 |
+
# nn.Conv2d(128, 128, kernel_size=3, padding=0),
|
33 |
+
nn.BatchNorm2d(128),
|
34 |
+
nn.ReLU(),
|
35 |
+
nn.MaxPool2d(2),
|
36 |
+
nn.Dropout(p=0.25)
|
37 |
+
)
|
38 |
+
|
39 |
+
self.conv_block_3 = nn.Sequential(
|
40 |
+
nn.Conv2d(128, 512, kernel_size=3, padding=1),
|
41 |
+
# nn.ReLU(),
|
42 |
+
# nn.Conv2d(128, 512, kernel_size=3, padding=0),
|
43 |
+
nn.BatchNorm2d(512),
|
44 |
+
nn.ReLU(),
|
45 |
+
nn.MaxPool2d(2),
|
46 |
+
nn.Dropout(p=0.25)
|
47 |
+
)
|
48 |
+
|
49 |
+
self.conv_block_4 = nn.Sequential(
|
50 |
+
nn.Conv2d(512, 512, kernel_size=3, padding=1),
|
51 |
+
# nn.ReLU(),
|
52 |
+
# nn.Conv2d(512, 512, kernel_size=3, padding=2),
|
53 |
+
nn.BatchNorm2d(512),
|
54 |
+
nn.ReLU(),
|
55 |
+
nn.MaxPool2d(2),
|
56 |
+
nn.Dropout(p=0.25)
|
57 |
+
)
|
58 |
+
|
59 |
+
self.fc_1 = nn.Sequential(
|
60 |
+
nn.Flatten(),
|
61 |
+
nn.Linear(in_features=256*392, out_features = 256),
|
62 |
+
nn.BatchNorm1d(256),
|
63 |
+
nn.ReLU(),
|
64 |
+
nn.Dropout(p=0.25)
|
65 |
+
)
|
66 |
+
|
67 |
+
self.fc_2 = nn.Sequential(
|
68 |
+
|
69 |
+
# Where did this in_features shape come from?
|
70 |
+
# It's because each layer of our network compresses and changes the shape of our inputs data.
|
71 |
+
nn.Linear(in_features=256,
|
72 |
+
out_features=512),
|
73 |
+
nn.BatchNorm1d(512),
|
74 |
+
nn.ReLU(),
|
75 |
+
nn.Dropout(p=0.25)
|
76 |
+
)
|
77 |
+
|
78 |
+
self.classifier = nn.Sequential(
|
79 |
+
nn.Linear(in_features=512,
|
80 |
+
out_features=output_shape)
|
81 |
+
)
|
82 |
+
|
83 |
+
def forward(self, x):
|
84 |
+
x = self.conv_block_1(x)
|
85 |
+
x = self.conv_block_2(x)
|
86 |
+
x = self.conv_block_3(x)
|
87 |
+
x = self.conv_block_4(x)
|
88 |
+
x = self.fc_1(x)
|
89 |
+
x = self.fc_2(x)
|
90 |
+
x = self.classifier(x)
|
91 |
+
return x
|
requirements.txt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
torch==2.2.0
|
2 |
+
torchvision==0.17.0
|
3 |
+
gradio==4.17.0
|