Update app.py

app.py

@@ -2,12 +2,13 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 import cv2
+from PIL import Image
 from huggingface_hub import from_pretrained_keras
 
 def resize_image(img_in,input_height,input_width):
     return cv2.resize( img_in, ( input_width,input_height) ,interpolation=cv2.INTER_NEAREST)
 
-def visualize_model_output(prediction):
+def visualize_model_output(prediction, img):
     unique_classes = np.unique(prediction[:,:,0])
     rgb_colors = {'0' : [0, 0, 0],
                  '1' : [255, 0, 0],
@@ -37,7 +38,13 @@ def visualize_model_output(prediction):
         output[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1]
         output[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2]
 
-    return output.astype(np.uint8)
+    output = output.astype(np.uint8)
+    im_pil_output = Image.fromarray(output)
+    im_pil = Image.fromarray(img)
+
+    im_pil.paste(im_pil_output, (0,0))
+
+    return im_pil.astype(np.uint8)
 
 
 def do_prediction(model_name, img):
@@ -169,7 +176,7 @@ def do_prediction(model_name, img):
             '''
             #prediction_true = prediction_true * -1
             #prediction_true = prediction_true + 1
-            return "No numerical output", visualize_model_output(prediction_true)
+            return "No numerical output", visualize_model_output(prediction_true,img)
         
         # catch-all (we should not reach this)
         case _: