SohomToom commited on
Commit
c778e32
·
verified ·
1 Parent(s): 89815d3
Files changed (1) hide show
  1. app.py +25 -1
app.py CHANGED
@@ -7,6 +7,17 @@ import numpy as np
7
  import gradio as gr
8
  from paddleocr import PaddleOCR
9
 
 
 
 
 
 
 
 
 
 
 
 
10
  ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='models/det', rec_model_dir='models/rec', cls_model_dir='models/cls')
11
 
12
  def classify_background_color(avg_color, white_thresh=230, black_thresh=50, yellow_thresh=100):
@@ -41,13 +52,25 @@ def sample_border_color(image, box, padding=2):
41
 
42
  def detect_text_boxes(image, max_dim=1280):
43
  try:
 
 
 
 
 
44
  # Resize large images to reduce memory load
45
  height, width = image.shape[:2]
46
  if max(height, width) > max_dim:
47
  scale = max_dim / float(max(height, width))
48
  image = cv2.resize(image, (int(width * scale), int(height * scale)))
49
 
50
- results = PaddleOCR.ocr(image, cls=True)
 
 
 
 
 
 
 
51
 
52
  if results is None or not results[0]:
53
  print("No OCR results found or OCR returned None.")
@@ -69,6 +92,7 @@ def detect_text_boxes(image, max_dim=1280):
69
  return []
70
 
71
 
 
72
  def remove_text_dynamic_fill(img_path, output_path):
73
  image = cv2.imread(img_path)
74
  if image is None:
 
7
  import gradio as gr
8
  from paddleocr import PaddleOCR
9
 
10
+ from PIL import Image
11
+
12
+ def is_valid_image(path):
13
+ try:
14
+ img = Image.open(path)
15
+ img.verify()
16
+ return True
17
+ except:
18
+ return False
19
+
20
+
21
  ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='models/det', rec_model_dir='models/rec', cls_model_dir='models/cls')
22
 
23
  def classify_background_color(avg_color, white_thresh=230, black_thresh=50, yellow_thresh=100):
 
52
 
53
  def detect_text_boxes(image, max_dim=1280):
54
  try:
55
+ # Check if image is valid
56
+ if image is None or not hasattr(image, 'shape'):
57
+ print("Invalid image. Skipping...")
58
+ return []
59
+
60
  # Resize large images to reduce memory load
61
  height, width = image.shape[:2]
62
  if max(height, width) > max_dim:
63
  scale = max_dim / float(max(height, width))
64
  image = cv2.resize(image, (int(width * scale), int(height * scale)))
65
 
66
+ # Ensure image is in RGB
67
+ if image.shape[2] == 1:
68
+ image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
69
+ elif image.shape[2] == 3:
70
+ image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
71
+
72
+ # Call PaddleOCR correctly
73
+ results = ocr.ocr(image, cls=True)
74
 
75
  if results is None or not results[0]:
76
  print("No OCR results found or OCR returned None.")
 
92
  return []
93
 
94
 
95
+
96
  def remove_text_dynamic_fill(img_path, output_path):
97
  image = cv2.imread(img_path)
98
  if image is None: