SWHL commited on
Commit
8095db7
1 Parent(s): ea0f4b5

fix(app): Remove the single visualize function

Browse files
Files changed (2) hide show
  1. app.py +6 -7
  2. utils.py +0 -78
app.py CHANGED
@@ -9,10 +9,10 @@ import numpy as np
9
  import pandas as pd
10
  import streamlit as st
11
  from PIL import Image
12
- from rapidocr_onnxruntime import RapidOCR
13
  from streamlit_image_select import image_select
14
 
15
- from utils import visualize
16
 
17
  font_dict = {
18
  "ch": "chinese_cht.ttf",
@@ -115,9 +115,8 @@ def inference(
115
  dt_boxes, rec_res, scores = list(zip(*ocr_result))
116
  font_path = Path("fonts") / font_dict.get(lan_name)
117
 
118
- vis_img = visualize(
119
- Image.fromarray(img), dt_boxes, rec_res, scores, font_path=str(font_path)
120
- )
121
  out_df = pd.DataFrame(
122
  [[rec, score] for rec, score in zip(rec_res, scores)],
123
  columns=("Rec", "Score"),
@@ -137,8 +136,8 @@ if __name__ == "__main__":
137
  )
138
  st.markdown(
139
  """
140
- <p align="left">
141
- <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>
142
  <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
143
  <a href="https://pepy.tech/project/rapidocr_onnxruntime"><img src="https://static.pepy.tech/personalized-badge/rapidocr_onnxruntime?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20Ort"></a>
144
  <a href="https://pypi.org/project/rapidocr-onnxruntime/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapidocr-onnxruntime"></a>
 
9
  import pandas as pd
10
  import streamlit as st
11
  from PIL import Image
12
+ from rapidocr_onnxruntime import RapidOCR, VisRes
13
  from streamlit_image_select import image_select
14
 
15
+ vis = VisRes()
16
 
17
  font_dict = {
18
  "ch": "chinese_cht.ttf",
 
115
  dt_boxes, rec_res, scores = list(zip(*ocr_result))
116
  font_path = Path("fonts") / font_dict.get(lan_name)
117
 
118
+ vis_img = vis(img, dt_boxes, rec_res, scores, font_path=str(font_path))
119
+ vis_img = vis_img[..., ::-1]
 
120
  out_df = pd.DataFrame(
121
  [[rec, score] for rec, score in zip(rec_res, scores)],
122
  columns=("Rec", "Score"),
 
136
  )
137
  st.markdown(
138
  """
139
+ <p align="center">
140
+ <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.13-aff.svg"></a>
141
  <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
142
  <a href="https://pepy.tech/project/rapidocr_onnxruntime"><img src="https://static.pepy.tech/personalized-badge/rapidocr_onnxruntime?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20Ort"></a>
143
  <a href="https://pypi.org/project/rapidocr-onnxruntime/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapidocr-onnxruntime"></a>
utils.py DELETED
@@ -1,78 +0,0 @@
1
- # -*- encoding: utf-8 -*-
2
- # @Author: SWHL
3
- # @Contact: liekkaskono@163.com
4
- import math
5
- import random
6
- from pathlib import Path
7
-
8
- import numpy as np
9
- from PIL import Image, ImageDraw, ImageFont
10
-
11
-
12
- def draw_ocr_box_txt(image, boxes, txts, font_path, scores=None, text_score=0.5):
13
- h, w = image.height, image.width
14
- img_left = image.copy()
15
- img_right = Image.new("RGB", (w, h), (255, 255, 255))
16
-
17
- random.seed(0)
18
- draw_left = ImageDraw.Draw(img_left)
19
- draw_right = ImageDraw.Draw(img_right)
20
- for idx, (box, txt) in enumerate(zip(boxes, txts)):
21
- if scores is not None and float(scores[idx]) < text_score:
22
- continue
23
-
24
- color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
25
-
26
- box = [tuple(v) for v in box]
27
- draw_left.polygon(box, fill=color)
28
- draw_right.text([box[3][0], box[3][1]], str(idx), fill=color)
29
-
30
- draw_right.polygon(
31
- [
32
- box[0][0],
33
- box[0][1],
34
- box[1][0],
35
- box[1][1],
36
- box[2][0],
37
- box[2][1],
38
- box[3][0],
39
- box[3][1],
40
- ],
41
- outline=color,
42
- )
43
-
44
- box_height = math.sqrt(
45
- (box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2
46
- )
47
-
48
- box_width = math.sqrt(
49
- (box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2
50
- )
51
-
52
- if box_height > 2 * box_width:
53
- font_size = max(int(box_width * 0.9), 10)
54
- font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
55
- cur_y = box[0][1]
56
- for c in txt:
57
- char_size = font.getsize(c)
58
- draw_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
59
- cur_y += char_size[1]
60
- else:
61
- font_size = max(int(box_height * 0.8), 10)
62
- font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
63
- draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
64
-
65
- img_left = Image.blend(image, img_left, 0.5)
66
- img_show = Image.new("RGB", (w * 2, h), (255, 255, 255))
67
- img_show.paste(img_left, (0, 0, w, h))
68
- img_show.paste(img_right, (w, 0, w * 2, h))
69
- return np.array(img_show)
70
-
71
-
72
- def visualize(image, boxes, txts, scores, font_path="./fonts/FZYTK.TTF"):
73
- draw_img = draw_ocr_box_txt(image, boxes, txts, font_path, scores, text_score=0.5)
74
-
75
- draw_img_save = Path("./inference_results/")
76
- if not draw_img_save.exists():
77
- draw_img_save.mkdir(parents=True, exist_ok=True)
78
- return draw_img[:, :, ::-1]