Update app.py

app.py

@@ -6,28 +6,24 @@ 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)
 
 
-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")
-
+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):
@@ -38,12 +34,10 @@ def process_data(content):
 
 
 def predict(img, size):
-    with tf.device(DEVICE):
-        img = tf.image.resize_with_pad(img, size, size)
-        img = tf.image.per_image_standardization(img)
-        data = process_data(image)
-        data = tf.expand_dims(data, 0)
-        out = model(data)[0]
+    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)))