Refactor sketch recognition app: update OpenCV dependency to headless version and simplify prediction function
Browse files- app.py +10 -29
- mnist-classes.png +0 -0
- requirements.txt +1 -1
app.py
CHANGED
@@ -1,10 +1,7 @@
|
|
1 |
-
import
|
2 |
-
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" # Disable oneDNN optimizations
|
3 |
-
|
4 |
import gradio as gr
|
5 |
import tensorflow as tf
|
6 |
import cv2
|
7 |
-
import numpy as np
|
8 |
|
9 |
# app title
|
10 |
title = "Welcome on your first sketch recognition app!"
|
@@ -12,7 +9,7 @@ title = "Welcome on your first sketch recognition app!"
|
|
12 |
# app description
|
13 |
head = (
|
14 |
"<center>"
|
15 |
-
"<img src='
|
16 |
"The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
|
17 |
"</center>"
|
18 |
)
|
@@ -29,34 +26,18 @@ labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight"
|
|
29 |
# load model (trained on MNIST dataset)
|
30 |
model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
|
31 |
|
32 |
-
#
|
33 |
def predict(img):
|
34 |
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
# Ensure grayscale format (convert from RGB if necessary)
|
40 |
-
if len(img.shape) == 3:
|
41 |
-
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
42 |
-
|
43 |
-
# Resize the image to 28x28
|
44 |
-
img = cv2.resize(img, (img_size, img_size))
|
45 |
-
|
46 |
-
# Normalize pixel values to [0, 1]
|
47 |
-
img = img / 255.0
|
48 |
-
|
49 |
-
# Reshape to match the model input shape
|
50 |
-
img = img.reshape(1, img_size, img_size, 1)
|
51 |
-
|
52 |
-
# Model predictions
|
53 |
-
preds = model.predict(img)[0]
|
54 |
|
55 |
-
|
56 |
-
|
57 |
|
58 |
-
|
59 |
-
|
60 |
|
61 |
# top 3 of classes
|
62 |
label = gr.Label(num_top_classes=3)
|
|
|
1 |
+
# import dependencies
|
|
|
|
|
2 |
import gradio as gr
|
3 |
import tensorflow as tf
|
4 |
import cv2
|
|
|
5 |
|
6 |
# app title
|
7 |
title = "Welcome on your first sketch recognition app!"
|
|
|
9 |
# app description
|
10 |
head = (
|
11 |
"<center>"
|
12 |
+
"<img src='./mnist-classes.png' width=400>"
|
13 |
"The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
|
14 |
"</center>"
|
15 |
)
|
|
|
26 |
# load model (trained on MNIST dataset)
|
27 |
model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
|
28 |
|
29 |
+
# prediction function for sketch recognition
|
30 |
def predict(img):
|
31 |
|
32 |
+
# image shape: 28x28x1
|
33 |
+
img = cv2.resize(img, (img_size, img_size))
|
34 |
+
img = img.reshape(1, img_size, img_size, 1)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
35 |
|
36 |
+
# model predictions
|
37 |
+
preds = model.predict(img)[0]
|
38 |
|
39 |
+
# return the probability for each classe
|
40 |
+
return {label: float(pred) for label, pred in zip(labels, preds)}
|
41 |
|
42 |
# top 3 of classes
|
43 |
label = gr.Label(num_top_classes=3)
|
mnist-classes.png
ADDED
requirements.txt
CHANGED
@@ -1,3 +1,3 @@
|
|
1 |
tensorflow
|
2 |
-
opencv-python
|
3 |
numpy
|
|
|
1 |
tensorflow
|
2 |
+
opencv-python-headless
|
3 |
numpy
|