Update app.py

app.py

@@ -211,15 +211,12 @@ def sepia(input_img):
     img_batch = np.expand_dims(img, axis=0)
     img_batch = np.transpose(img_batch, (0, 3, 1, 2))
     
-    outputs = sess.run(None, {"pixel_values": img_batch})[0]    
-    logits = outputs.logits
-
-    logits = tf.transpose(logits, [0, 2, 3, 1])
-    logits = tf.image.resize(
-        logits, input_img.size[::-1]
-    )  # We reverse the shape of `image` because `image.size` returns width and height.
-    seg = tf.math.argmax(logits, axis=-1)[0]
+    logits = sess.run(None, {"pixel_values": img_batch})[0]
 
+    logits = np.transpose(logits, (0, 2, 3, 1))
+    seg = np.argmax(logits, axis=-1)[0].astype('float32')
+    seg = cv2.resize(seg, (640, 640))
+    
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )  # height, width, 3