Update app.py

app.py

@@ -182,14 +182,12 @@ def do_prediction(model_name, img):
             return "Found {} columns".format(num_col), None
 
         case "SBB/eynollah-page-extraction":
-
             img_height_model = model.layers[len(model.layers) - 1].output_shape[1]
             img_width_model = model.layers[len(model.layers) - 1].output_shape[2]
     
             img_h_page = img.shape[0]
             img_w_page = img.shape[1]
 
-            print(img.shape, img_h_page,img_w_page)
             img = img / float(255.0)
             img = resize_image(img, img_height_model, img_width_model)
 
@@ -206,7 +204,6 @@ def do_prediction(model_name, img):
             #thresh = cv2.dilate(thresh, KERNEL, iterations=3)
             contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
 
-            print(len(contours),'ggg')
             if len(contours)>0:
                 cnt_size = np.array([cv2.contourArea(contours[j]) for j in range(len(contours))])
                 cnt = contours[np.argmax(cnt_size)]
@@ -224,7 +221,6 @@ def do_prediction(model_name, img):
 
                 box = [x, y, w, h]
 
-                print(box)
 
                 img_border = np.zeros((prediction_true.shape[0],prediction_true.shape[1]))
                 img_border[y:y+h, x:x+w] = 1
@@ -308,9 +304,12 @@ def do_prediction(model_name, img):
                     label_p_pred = model.predict(img_patch.reshape(1, img_patch.shape[0], img_patch.shape[1], img_patch.shape[2]),
                                                  verbose=0)
 
-                    seg = np.argmax(label_p_pred, axis=3)[0]
-
-                    seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)
+                    if model_name == "SBB/eynollah-enhancement":
+                        seg_color = label_p_pred[0, :, :, :]
+                        seg_color = seg_color * 255
+                    else:     
+                        seg = np.argmax(label_p_pred, axis=3)[0]
+                        seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)
 
                     if i == 0 and j == 0:
                         seg_color = seg_color[0 : seg_color.shape[0] - margin, 0 : seg_color.shape[1] - margin, :]
@@ -358,9 +357,14 @@ def do_prediction(model_name, img):
                         #mask_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - margin] = seg
                         prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - margin, :] = seg_color
 
-            prediction_true = prediction_true.astype(np.uint8)            
+            if model_name == "SBB/eynollah-enhancement":
+                prediction_true = prediction_true.astype(int)
+                return "No numerical output", prediction_true
+            else:
+                prediction_true = prediction_true.astype(np.uint8)
+                return "No numerical output", visualize_model_output(prediction_true,img_org, model_name)
 
-            return "No numerical output", visualize_model_output(prediction_true,img_org, model_name)
+            
         
         # catch-all (we should not reach this)
         case _: