Sushrut98 commited on
Commit
05101d2
·
verified ·
1 Parent(s): 9376853

Upload 5 files

Browse files
Files changed (5) hide show
  1. app.py +90 -0
  2. class_names.txt +9 -0
  3. cnn.pth +3 -0
  4. model.py +91 -0
  5. 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