AnkitShrestha commited on
Commit
f2acee5
·
1 Parent(s): 6d7bbfa

minor fixes

Browse files
Files changed (1) hide show
  1. utils.py +5 -5
utils.py CHANGED
@@ -1,7 +1,7 @@
1
  from doctr.models import detection_predictor, recognition_predictor
2
  from doctr.io import DocumentFile
3
- from surya.recognition import RecognitionPredictor as SuryaRecognitionPredictor
4
- from surya.detection import DetectionPredictor as SuryaDetectionPredictor
5
  from PIL import Image
6
  # from functools import lru_cache
7
  from torchvision import models
@@ -253,9 +253,9 @@ def initialize_detector():
253
  if doctr_detector is None:
254
  doctr_detector = detection_predictor('db_mobilenet_v3_large', pretrained=True, assume_straight_pages=True, preserve_aspect_ratio=True)
255
  if surya_recognition_predictor is None:
256
- surya_recognition_predictor = SuryaRecognitionPredictor()
257
  if surya_detection_predictor is None:
258
- surya_detection_predictor = SuryaDetectionPredictor()
259
  return doctr_detector, surya_recognition_predictor, surya_detection_predictor
260
 
261
  def get_cleaned_boxes(out, page):
@@ -347,7 +347,7 @@ def ocr_citizenship_utils(image_path: str) -> List[List[str]]:
347
 
348
  # OCR PART
349
  langs = ["en",'ne']
350
- predictions = surya_recognition_predictor([pil_image], [langs],surya_detection_predictor)
351
  text_combo = ''
352
  for text_line in predictions[0].text_lines:
353
  text_combo = text_combo + " " + text_line.text.strip()
 
1
  from doctr.models import detection_predictor, recognition_predictor
2
  from doctr.io import DocumentFile
3
+ from surya.recognition import RecognitionPredictor
4
+ from surya.detection import DetectionPredictor
5
  from PIL import Image
6
  # from functools import lru_cache
7
  from torchvision import models
 
253
  if doctr_detector is None:
254
  doctr_detector = detection_predictor('db_mobilenet_v3_large', pretrained=True, assume_straight_pages=True, preserve_aspect_ratio=True)
255
  if surya_recognition_predictor is None:
256
+ surya_recognition_predictor = RecognitionPredictor()
257
  if surya_detection_predictor is None:
258
+ surya_detection_predictor = DetectionPredictor()
259
  return doctr_detector, surya_recognition_predictor, surya_detection_predictor
260
 
261
  def get_cleaned_boxes(out, page):
 
347
 
348
  # OCR PART
349
  langs = ["en",'ne']
350
+ predictions = surya_recognition_predictor(images=[pil_image], langs=[langs],det_predictor=surya_detection_predictor)
351
  text_combo = ''
352
  for text_line in predictions[0].text_lines:
353
  text_combo = text_combo + " " + text_line.text.strip()