narugo1992 commited on
Commit
4f6e58b
1 Parent(s): d10a366

dev(narugo): add this example

Browse files
Files changed (2) hide show
  1. app.py +0 -1
  2. detect.py +22 -11
app.py CHANGED
@@ -7,7 +7,6 @@ from detect import _ALL_MODELS, _DEFAULT_MODEL, detect_text
7
 
8
 
9
  def _gr_detect_text(image, model: str, threshold: float):
10
- print(image)
11
  return detection_visualize(image, detect_text(image, model, threshold))
12
 
13
 
 
7
 
8
 
9
  def _gr_detect_text(image, model: str, threshold: float):
 
10
  return detection_visualize(image, detect_text(image, model, threshold))
11
 
12
 
detect.py CHANGED
@@ -1,5 +1,6 @@
1
  import os.path
2
  from functools import lru_cache
 
3
 
4
  import cv2
5
  import numpy as np
@@ -18,7 +19,7 @@ def _get_available_models():
18
 
19
 
20
  _ALL_MODELS = list(_get_available_models())
21
- _DEFAULT_MODEL = 'dbnetpp_resnet50-oclip_fpnc_1200e_icdar2015'
22
 
23
 
24
  @lru_cache()
@@ -29,7 +30,7 @@ def _get_onnx_session(model):
29
  ))
30
 
31
 
32
- def detect_text(image: ImageTyping, model: str = _DEFAULT_MODEL, threshold: float = 0.05):
33
  origin_width, origin_height = width, height = image.size
34
  align = 32
35
  if width % align != 0:
@@ -52,17 +53,27 @@ def detect_text(image: ImageTyping, model: str = _DEFAULT_MODEL, threshold: floa
52
  heatmap = output_[0]
53
  heatmap = heatmap[:origin_height, :origin_width]
54
 
55
- cnts = cv2.findContours((heatmap * 255.0).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
56
- cnts = cnts[0] if len(cnts) == 2 else cnts[1]
 
 
 
 
 
 
57
  bboxes = []
58
- for c in cnts:
59
  x, y, w, h = cv2.boundingRect(c)
60
- x0, y0 = x, y
61
- x1, y1 = x + w, y + h
62
- area = heatmap[y0:y1, x0:x1]
63
- valid_area = area[area >= 1e-4]
64
- score = valid_area.mean().item()
65
  if score >= threshold:
66
- bboxes.append(((x0, y0, x1, y1), 'text', score))
 
 
 
67
 
 
 
 
 
68
  return bboxes
 
1
  import os.path
2
  from functools import lru_cache
3
+ from typing import List, Tuple
4
 
5
  import cv2
6
  import numpy as np
 
19
 
20
 
21
  _ALL_MODELS = list(_get_available_models())
22
+ _DEFAULT_MODEL = 'dbnetpp_resnet50_fpnc_1200e_icdar2015'
23
 
24
 
25
  @lru_cache()
 
30
  ))
31
 
32
 
33
+ def _get_heatmap_of_text(image: ImageTyping, model: str) -> np.ndarray:
34
  origin_width, origin_height = width, height = image.size
35
  align = 32
36
  if width % align != 0:
 
53
  heatmap = output_[0]
54
  heatmap = heatmap[:origin_height, :origin_width]
55
 
56
+ return heatmap
57
+
58
+
59
+ def _get_bounding_box_of_text(image: ImageTyping, model: str, threshold: float) \
60
+ -> List[Tuple[Tuple[int, int, int, int], float]]:
61
+ heatmap = _get_heatmap_of_text(image, model)
62
+ c_rets = cv2.findContours((heatmap * 255.0).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
63
+ contours = c_rets[0] if len(c_rets) == 2 else c_rets[1]
64
  bboxes = []
65
+ for c in contours:
66
  x, y, w, h = cv2.boundingRect(c)
67
+ x0, y0, x1, y1 = x, y, x + w, y + h
68
+ score = heatmap[y0:y1, x0:x1].mean().item()
 
 
 
69
  if score >= threshold:
70
+ bboxes.append(((x0, y0, x1, y1), score))
71
+
72
+ return bboxes
73
+
74
 
75
+ def detect_text(image: ImageTyping, model: str = _DEFAULT_MODEL, threshold: float = 0.05):
76
+ bboxes = []
77
+ for (x0, y0, x1, y1), score in _get_bounding_box_of_text(image, model, threshold):
78
+ bboxes.append(((x0, y0, x1, y1), 'text', score))
79
  return bboxes