a

app.py

@@ -45,20 +45,21 @@ def greet(input_img):
     outputs = model(**inputs)
     logits = outputs.logits
 
-    # Use .detach().numpy() to convert PyTorch tensor to NumPy array
-    logits = logits.detach().numpy()
-    logits = np.transpose(logits, [0, 2, 3, 1])
+    logits_np = logits.detach().numpy()
 
-    logits = np.resize(logits, input_img.size[::-1])
+    logits_tf = tf.convert_to_tensor(logits_np)
+    logits_tf = tf.transpose(logits_tf, [0, 2, 3, 1])
 
-    seg = np.argmax(logits, axis=-1)[0]
+    logits_tf = tf.image.resize(
+        logits_tf, input_img.size[::-1]
+    )
+    seg = tf.math.argmax(logits_tf, axis=-1)[0]
 
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )
-
     for label, color in enumerate(colormap):
-        color_seg[seg == label, :] = color
+        color_seg[seg.numpy() == label, :] = color
 
     pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
     pred_img = pred_img.astype(np.uint8)