-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
66 lines (45 loc) · 1.5 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# Import functions to database
from models.Images import Images
# Import funtions to process images
import numpy as np
import cv2
from PIL import Image
from io import BytesIO
import base64
import pytesseract
# Import functions to API
from fastapi import FastAPI, File, UploadFile
app = FastAPI()
# Convert image to numpy array
def read_image(file):
image = Image.open(BytesIO(file))
return image
# Clean noises from image
def clean_image(image):
img = np.array(image)
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
denoise_2 = cv2.fastNlMeansDenoisingColored(img,None,5,5,7,21)
return denoise_2
# Return text extracted from image
def get_ocr_from_image(image):
pytesseract.pytesseract.tesseract_cmd = 'C:\\Program Files\\Tesseract-OCR\\tesseract'
result = pytesseract.image_to_string(image, lang="eng")
return result
@app.post("/images")
async def process_image(file: UploadFile = File(...)):
extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
if not extension:
return "Image must be jpg or png format!"
formatted_image = read_image(await file.read())
cleaned_image = clean_image(formatted_image)
text = get_ocr_from_image(cleaned_image)
encoded_image = base64.b64encode(cleaned_image)
img = Images(file.filename, text, encoded_image)
image_id = img.insert()
payload = {
"filename": file.filename,
"text": text,
"data": encoded_image,
"image_id": str(image_id)
}
return payload