a

app.py

@@ -45,10 +45,7 @@ def greet(input_img):
     outputs = model(**inputs)
     logits = outputs.logits
 
-    logits_np = logits.detach().numpy()
-
-    logits_tf = tf.convert_to_tensor(logits_np)
-    logits_tf = tf.transpose(logits_tf, [0, 2, 3, 1])
+    logits_tf = tf.transpose(logits, [0, 2, 3, 1])
 
     logits_tf = tf.image.resize(
         logits_tf, input_img.size[::-1]
@@ -57,7 +54,7 @@ def greet(input_img):
 
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )
+    )  # height, width, 3
     for label, color in enumerate(colormap):
         color_seg[seg.numpy() == label, :] = color
 
@@ -67,6 +64,7 @@ def greet(input_img):
     fig = draw_plot(pred_img, seg)
     return fig
 
+
 def draw_plot(pred_img, seg):
     fig = plt.figure(figsize=(20, 15))
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])