Update app.py

app.py

@@ -20,43 +20,48 @@ def otsu_copy_binary(img):
 
     return img_r
 
-def visualize_model_output(prediction, img):
-    unique_classes = np.unique(prediction[:,:,0])
-    rgb_colors = {'0' : [255, 255, 255],
-                 '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(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]
-
-
-
-    img = resize_image(img, output.shape[0], output.shape[1])
-
-    output = output.astype(np.int32)
-    img = img.astype(np.int32)
-
+def visualize_model_output(prediction, img, model_name):
+    if model_name == "SBB/eynollah-binarization":
+        prediction_true = prediction_true * -1
+        prediction_true = prediction_true + 1
+        added_image = prediction_true * 255
+    else:
+        unique_classes = np.unique(prediction[:,:,0])
+        rgb_colors = {'0' : [255, 255, 255],
+                     '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(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]
+    
     
     
-    added_image = cv2.addWeighted(img,0.5,output,0.1,0)
+        img = resize_image(img, output.shape[0], output.shape[1])
+    
+        output = output.astype(np.int32)
+        img = img.astype(np.int32)
+    
+        
+        
+        added_image = cv2.addWeighted(img,0.5,output,0.1,0)
     return added_image
 
 def return_num_columns(img):
@@ -75,7 +80,7 @@ def return_num_columns(img):
     return num_col
 
 def return_scaled_image(img, num_col, width_early, model_name):
-    if model_name == "SBB/eynollah-main-regions-aug-rotation" or "SBB/eynollah-main-regions-aug-scaling" or "SBB/eynollah-main-regions-ensembled" or "SBB/eynollah-textline":
+    if model_name == "SBB/eynollah-main-regions-aug-rotation" or "SBB/eynollah-main-regions-aug-scaling" or "SBB/eynollah-main-regions-ensembled" or "SBB/eynollah-textline" or "SBB/eynollah-binarization":
         if num_col == 1 and width_early < 1100:
             img_w_new = 2000
             img_h_new = int(img.shape[0] / float(img.shape[1]) * 2000)
@@ -313,7 +318,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,img_org)
+            return "No numerical output", visualize_model_output(prediction_true,img_org, model_name)
         
         # catch-all (we should not reach this)
         case _: