adding visualizer

app.py

@@ -7,6 +7,36 @@ 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):
+    unique_classes = np.unique(prediction[:,:,0])
+    rgb_colors = {'0':[0, 0, 0],
+                 '1', [255, 0, 0],
+                 '2', [255, 125, 0],
+                 '3', [255, 0, 125],
+                 '4', [125, 125, 125],
+                 '5', [125, 125, 0],
+                 '6', [0, 125, 255],
+                 '7', [0, 125, 0],
+                 '8', [125, 125, 125],
+                 '9', [0, 125, 255],
+                 '10', [125, 0, 125],
+                 '11', [0, 255, 0],
+                 '12', [0, 0, 255],
+                 '13', [0, 255, 255],
+                 '14', [255, 125, 125],
+                 '15', [255, 0, 255]}
+
+    output = np.zeros(prediction.shape)
+
+    for unq_class in unique_classes:
+        rgb_class_unique = rgb_colors[str(unq_class)]
+        output[:,:,0] = rgb_class_unique[0]
+        output[:,:,1] = rgb_class_unique[1]
+        output[:,:,2] = rgb_class_unique[2]
+
+    return output
+
+
 def do_prediction(model_name, img):
     model = from_pretrained_keras(model_name)
 
@@ -136,7 +166,7 @@ def do_prediction(model_name, img):
             '''
             #prediction_true = prediction_true * -1
             #prediction_true = prediction_true + 1
-            return "No numerical output", prediction_true
+            return "No numerical output", visualize_model_output(prediction_true)
         
         # catch-all (we should not reach this)
         case _: