vahidrezanezhad commited on
Commit
437292f
·
verified ·
1 Parent(s): acb5dba

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -37
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
- unique_classes = np.unique(prediction[:,:,0])
25
- rgb_colors = {'0' : [255, 255, 255],
26
- '1' : [255, 0, 0],
27
- '2' : [255, 125, 0],
28
- '3' : [255, 0, 125],
29
- '4' : [125, 125, 125],
30
- '5' : [125, 125, 0],
31
- '6' : [0, 125, 255],
32
- '7' : [0, 125, 0],
33
- '8' : [125, 125, 125],
34
- '9' : [0, 125, 255],
35
- '10' : [125, 0, 125],
36
- '11' : [0, 255, 0],
37
- '12' : [0, 0, 255],
38
- '13' : [0, 255, 255],
39
- '14' : [255, 125, 125],
40
- '15' : [255, 0, 255]}
41
-
42
- output = np.zeros(prediction.shape)
43
-
44
- for unq_class in unique_classes:
45
- rgb_class_unique = rgb_colors[str(int(unq_class))]
46
- output[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0]
47
- output[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1]
48
- output[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2]
49
-
50
-
51
-
52
- img = resize_image(img, output.shape[0], output.shape[1])
53
-
54
- output = output.astype(np.int32)
55
- img = img.astype(np.int32)
56
-
57
 
58
 
59
- added_image = cv2.addWeighted(img,0.5,output,0.1,0)
 
 
 
 
 
 
 
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 _: