Update app.py

app.py

@@ -31,10 +31,11 @@ def visualize_model_output(prediction):
     print(output.shape,'shape of output')
 
     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]
+        print(unq_class,'unq_class')
+        rgb_class_unique = rgb_colors[str(int(unq_class))]
+        output[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0]
+        output[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1]
+        output[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2]
 
     return output