Update app.py

app.py

@@ -6,44 +6,45 @@ import tensorflow.keras as keras
 from gradio import inputs, outputs
 
 SIZE = 256
+DEVICE = "/cpu:0"
+
 
 with open("./tags.json", "rt", encoding="utf-8") as f:
     tags = json.load(f)
 
 
-base_model = keras.applications.resnet.ResNet50(
-    include_top=False, weights=None, input_shape=(SIZE, SIZE, 3)
-)
-model = keras.Sequential(
-    [
-        base_model,
-        keras.layers.Conv2D(filters=len(tags), kernel_size=(1, 1), padding="same"),
-        keras.layers.BatchNormalization(epsilon=1.001e-5),
-        keras.layers.GlobalAveragePooling2D(name="avg_pool"),
-        keras.layers.Activation("sigmoid"),
-    ]
-)
-model.load_weights("tf_model.h5")
-
-@tf.function
-def process_data(content):
-    img = tf.io.decode_jpeg(content, channels=3)
-    img = tf.image.resize_with_pad(img, SIZE, SIZE)
-    img = tf.image.per_image_standardization(img)
-    return img
-
-
-def predict(img, size):
-    img = tf.image.resize_with_pad(img, size, size)
-    img = tf.image.per_image_standardization(img)
-    data = tf.expand_dims(img, 0)
-    out,*_ = model(data)
-    return dict((tags[i], out[i].numpy()) for i in range(len(tags)))
+with tf.device(DEVICE):
+    base_model = keras.applications.resnet.ResNet50(
+        include_top=False, weights=None, input_shape=(SIZE, SIZE, 3)
+    )
+    model = keras.Sequential(
+        [
+            base_model,
+            keras.layers.Conv2D(filters=len(tags), kernel_size=(1, 1), padding="same"),
+            keras.layers.BatchNormalization(epsilon=1.001e-5),
+            keras.layers.GlobalAveragePooling2D(name="avg_pool"),
+            keras.layers.Activation("sigmoid"),
+        ]
+    )
+    model.load_weights("tf_model.h5")
+
+
+def predict(img, hide: float):
+    with tf.device(DEVICE):
+        img = tf.image.resize_with_pad(img, SIZE, SIZE)
+        img = tf.image.per_image_standardization(img)
+        data = tf.expand_dims(img, 0)
+        out, *_ = model(data)
+    labels = {tag: float(out[i].numpy()) for i, tag in enumerate(tags)}
+    return {k: v for k, v in labels.items() if v >= hide}
 
 
 image = inputs.Image(label="Upload your image here!")
-size = inputs.Number(label="Image resize", default=SIZE)
+hide_threshold = inputs.Slider(
+    label="Hide confidence lower than", default=0.5, maximum=1, minimum=0
+)
 
-labels = outputs.Label(label="Tags")
+labels = outputs.Label(label="Tags", type="confidences")
 
-gr.Interface(predict, inputs=[image, size], outputs=[labels])
+interface = gr.Interface(predict, inputs=[image, hide_threshold], outputs=[labels])
+interface.launch()