coutant commited on
Commit
b371285
1 Parent(s): 3d3c04e

an option to detection with craft_hw_ocr versus craft_text_detector

Browse files
Files changed (2) hide show
  1. app.py +17 -9
  2. requirements.txt +1 -0
app.py CHANGED
@@ -2,9 +2,6 @@ import PIL.Image
2
  import gradio as gr
3
  import torch
4
  import numpy as np
5
- from craft_text_detector import Craft
6
-
7
- craft = Craft(output_dir='output', crop_type="box", cuda=torch.cuda.is_available(), export_extra=True)
8
 
9
  dw=0.3
10
  dh=0.25
@@ -45,20 +42,31 @@ def is_signature(prediction_result) -> bool:
45
  return True
46
  return False
47
 
48
- def detect(image: PIL.Image.Image):
 
 
49
  result = craft.detect_text( np.asarray(image))
50
- return result['boxes'], is_signature(result)
 
 
 
 
 
 
 
 
 
 
51
 
52
- def process(image:PIL.Image.Image):
53
  if image is None:
54
  return None,0,False
55
- boxes,signed = detect( image)
56
- annotated = PIL.Image.open('output/image_text_detection.png') # image with boxes displayed
57
  return annotated, len(boxes), signed
58
 
59
  gr.Interface(
60
  fn = process,
61
- inputs = [ gr.Image(type="pil", label="Input") ],
62
  outputs = [ gr.Image(type="pil", label="Output"), gr.Label(label="nb of text detections"), gr.Label(label="Has signature") ],
63
  title="Detect signature in image",
64
  description="Is the photo or image watermarked by a signature?",
 
2
  import gradio as gr
3
  import torch
4
  import numpy as np
 
 
 
5
 
6
  dw=0.3
7
  dh=0.25
 
42
  return True
43
  return False
44
 
45
+ def detect_with_craft_text_detector(image: PIL.Image.Image):
46
+ from craft_text_detector import Craft
47
+ craft = Craft(output_dir='output', crop_type="box", cuda=torch.cuda.is_available(), export_extra=True)
48
  result = craft.detect_text( np.asarray(image))
49
+ annotated = PIL.Image.open('output/image_text_detection.png') # image with boxes displayed
50
+ return annotated, result['boxes'], is_signature(result)
51
+
52
+ def detect_with_craft_hw_ocr(image: PIL.Image.Image):
53
+ image = np.asarray(image)
54
+ from craft_hw_ocr import OCR
55
+ ocr = OCR.load_models()
56
+ image, results = OCR.detection(image, ocr[2])
57
+ bboxes, _ = OCR.recoginition(image, results, ocr[0], ocr[1])
58
+ annotated = OCR.visualize(image, results)
59
+ return annotated, bboxes, False
60
 
61
+ def process(image:PIL.Image.Image, lib:str):
62
  if image is None:
63
  return None,0,False
64
+ annotated, boxes, signed = detect_with_craft_text_detector(image) if lib=='craft_text_detector' else detect_with_craft_hw_ocr( image)
 
65
  return annotated, len(boxes), signed
66
 
67
  gr.Interface(
68
  fn = process,
69
+ inputs = [ gr.Image(type="pil", label="Input"), gr.inputs.Radio(label='Model', choices=["craft_text_detector", "craft_hw_ocr"], default='craft_text_detector') ],
70
  outputs = [ gr.Image(type="pil", label="Output"), gr.Label(label="nb of text detections"), gr.Label(label="Has signature") ],
71
  title="Detect signature in image",
72
  description="Is the photo or image watermarked by a signature?",
requirements.txt CHANGED
@@ -9,3 +9,4 @@ matplotlib
9
  scipy
10
  psutil
11
  craft_text_detector
 
 
9
  scipy
10
  psutil
11
  craft_text_detector
12
+ craft_hw_ocr