Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -20,43 +20,48 @@ def otsu_copy_binary(img):
|
|
20 |
|
21 |
return img_r
|
22 |
|
23 |
-
def visualize_model_output(prediction, img):
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
output
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
img = img.astype(np.int32)
|
56 |
-
|
57 |
|
58 |
|
59 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
60 |
return added_image
|
61 |
|
62 |
def return_num_columns(img):
|
@@ -75,7 +80,7 @@ def return_num_columns(img):
|
|
75 |
return num_col
|
76 |
|
77 |
def return_scaled_image(img, num_col, width_early, model_name):
|
78 |
-
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":
|
79 |
if num_col == 1 and width_early < 1100:
|
80 |
img_w_new = 2000
|
81 |
img_h_new = int(img.shape[0] / float(img.shape[1]) * 2000)
|
@@ -313,7 +318,7 @@ def do_prediction(model_name, img):
|
|
313 |
'''
|
314 |
#prediction_true = prediction_true * -1
|
315 |
#prediction_true = prediction_true + 1
|
316 |
-
return "No numerical output", visualize_model_output(prediction_true,img_org)
|
317 |
|
318 |
# catch-all (we should not reach this)
|
319 |
case _:
|
|
|
20 |
|
21 |
return img_r
|
22 |
|
23 |
+
def visualize_model_output(prediction, img, model_name):
|
24 |
+
if model_name == "SBB/eynollah-binarization":
|
25 |
+
prediction_true = prediction_true * -1
|
26 |
+
prediction_true = prediction_true + 1
|
27 |
+
added_image = prediction_true * 255
|
28 |
+
else:
|
29 |
+
unique_classes = np.unique(prediction[:,:,0])
|
30 |
+
rgb_colors = {'0' : [255, 255, 255],
|
31 |
+
'1' : [255, 0, 0],
|
32 |
+
'2' : [255, 125, 0],
|
33 |
+
'3' : [255, 0, 125],
|
34 |
+
'4' : [125, 125, 125],
|
35 |
+
'5' : [125, 125, 0],
|
36 |
+
'6' : [0, 125, 255],
|
37 |
+
'7' : [0, 125, 0],
|
38 |
+
'8' : [125, 125, 125],
|
39 |
+
'9' : [0, 125, 255],
|
40 |
+
'10' : [125, 0, 125],
|
41 |
+
'11' : [0, 255, 0],
|
42 |
+
'12' : [0, 0, 255],
|
43 |
+
'13' : [0, 255, 255],
|
44 |
+
'14' : [255, 125, 125],
|
45 |
+
'15' : [255, 0, 255]}
|
46 |
+
|
47 |
+
output = np.zeros(prediction.shape)
|
48 |
+
|
49 |
+
for unq_class in unique_classes:
|
50 |
+
rgb_class_unique = rgb_colors[str(int(unq_class))]
|
51 |
+
output[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0]
|
52 |
+
output[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1]
|
53 |
+
output[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2]
|
54 |
+
|
|
|
|
|
55 |
|
56 |
|
57 |
+
img = resize_image(img, output.shape[0], output.shape[1])
|
58 |
+
|
59 |
+
output = output.astype(np.int32)
|
60 |
+
img = img.astype(np.int32)
|
61 |
+
|
62 |
+
|
63 |
+
|
64 |
+
added_image = cv2.addWeighted(img,0.5,output,0.1,0)
|
65 |
return added_image
|
66 |
|
67 |
def return_num_columns(img):
|
|
|
80 |
return num_col
|
81 |
|
82 |
def return_scaled_image(img, num_col, width_early, model_name):
|
83 |
+
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":
|
84 |
if num_col == 1 and width_early < 1100:
|
85 |
img_w_new = 2000
|
86 |
img_h_new = int(img.shape[0] / float(img.shape[1]) * 2000)
|
|
|
318 |
'''
|
319 |
#prediction_true = prediction_true * -1
|
320 |
#prediction_true = prediction_true + 1
|
321 |
+
return "No numerical output", visualize_model_output(prediction_true,img_org, model_name)
|
322 |
|
323 |
# catch-all (we should not reach this)
|
324 |
case _:
|