kavg commited on
Commit
16d7e9b
·
1 Parent(s): ccb5ac8

Revert "commit before changing entity merging process"

Browse files
Files changed (3) hide show
  1. main.py +10 -3
  2. ocr.py +4 -4
  3. preprocess.py +0 -9
main.py CHANGED
@@ -103,12 +103,19 @@ def ApplyOCR(content):
103
  try:
104
  trocr_client = ocr.TrOCRClient(config['settings'].TROCR_API_URL)
105
  handwritten_ocr_df = trocr_client.ocr(handwritten_imgs, image)
106
- except Exception as e:
107
- print(e)
108
  raise HTTPException(status_code=400, detail="handwritten OCR process failed")
 
 
 
 
 
 
 
 
 
109
 
110
  ocr_df = pd.concat([handwritten_ocr_df, printed_ocr_df])
111
- # ocr_df = printed_ocr_df
112
  return ocr_df, image
113
 
114
 
 
103
  try:
104
  trocr_client = ocr.TrOCRClient(config['settings'].TROCR_API_URL)
105
  handwritten_ocr_df = trocr_client.ocr(handwritten_imgs, image)
106
+ except:
 
107
  raise HTTPException(status_code=400, detail="handwritten OCR process failed")
108
+
109
+ try:
110
+ jpeg_bytes = io.BytesIO()
111
+ printed_img.save(jpeg_bytes, format='JPEG')
112
+ jpeg_content = jpeg_bytes.getvalue()
113
+ vision_client = ocr.VisionClient(config['settings'].GCV_AUTH)
114
+ printed_ocr_df = vision_client.ocr(jpeg_content, printed_img)
115
+ except:
116
+ raise HTTPException(status_code=400, detail="Printed OCR process failed")
117
 
118
  ocr_df = pd.concat([handwritten_ocr_df, printed_ocr_df])
 
119
  return ocr_df, image
120
 
121
 
ocr.py CHANGED
@@ -1,11 +1,12 @@
1
  from google.cloud import vision
2
  from google.oauth2 import service_account
 
3
  import pandas as pd
4
  import json
5
  import numpy as np
 
6
  import io
7
  import requests
8
- from preprocess import cam_scanner_filter
9
 
10
  image_ext = ("*.jpg", "*.jpeg", "*.png")
11
 
@@ -22,7 +23,7 @@ class VisionClient:
22
  except ValueError as e:
23
  print("Image could not be read")
24
  return
25
- response = self.client.document_text_detection(image, timeout=60)
26
  return response
27
 
28
  def get_response(self, content):
@@ -133,8 +134,7 @@ class TrOCRClient():
133
  boxObjects = []
134
  for i in range(len(handwritten_imgs)):
135
  handwritten_img = handwritten_imgs[i]
136
- handwritten_img_processed = cam_scanner_filter(handwritten_img[0])
137
- ocr_result = self.send_request(handwritten_img_processed)
138
  boxObjects.append({
139
  "id": i-1,
140
  "text": ocr_result,
 
1
  from google.cloud import vision
2
  from google.oauth2 import service_account
3
+ from google.protobuf.json_format import MessageToJson
4
  import pandas as pd
5
  import json
6
  import numpy as np
7
+ from PIL import Image
8
  import io
9
  import requests
 
10
 
11
  image_ext = ("*.jpg", "*.jpeg", "*.png")
12
 
 
23
  except ValueError as e:
24
  print("Image could not be read")
25
  return
26
+ response = self.client.document_text_detection(image, timeout=10)
27
  return response
28
 
29
  def get_response(self, content):
 
134
  boxObjects = []
135
  for i in range(len(handwritten_imgs)):
136
  handwritten_img = handwritten_imgs[i]
137
+ ocr_result = self.send_request(handwritten_img[0])
 
138
  boxObjects.append({
139
  "id": i-1,
140
  "text": ocr_result,
preprocess.py CHANGED
@@ -1,8 +1,5 @@
1
  import torch
2
  from transformers import AutoTokenizer
3
- import cv2
4
- from PIL import Image
5
- import numpy as np
6
 
7
  def normalize_box(box, width, height):
8
  return [
@@ -12,12 +9,6 @@ def normalize_box(box, width, height):
12
  int(1000 * (box[3] / height)),
13
  ]
14
 
15
- def cam_scanner_filter(img):
16
- image1 = np.array(img)
17
- img = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
18
- thresh2 = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY, 199, 15)
19
- return Image.fromarray(thresh2)
20
-
21
  # class to turn the keys of a dict into attributes (thanks Stackoverflow)
22
  class AttrDict(dict):
23
  def __init__(self, *args, **kwargs):
 
1
  import torch
2
  from transformers import AutoTokenizer
 
 
 
3
 
4
  def normalize_box(box, width, height):
5
  return [
 
9
  int(1000 * (box[3] / height)),
10
  ]
11
 
 
 
 
 
 
 
12
  # class to turn the keys of a dict into attributes (thanks Stackoverflow)
13
  class AttrDict(dict):
14
  def __init__(self, *args, **kwargs):