Binarization in fastapi side

binarization .py

@@ -0,0 +1,107 @@
+import math
+import cv2
+from PIL import Image
+import numpy as np
+
+def integral(img):
+    '''
+    计算图像的积分和平方积分
+    :param img:Mat--- 输入待处理图像
+    :return:integral_sum, integral_sqrt_sum：Mat--- 积分图和平方积分图
+    '''
+    integral_sum=np.zeros((img.shape[0],img.shape[1]),dtype=np.int32)
+    integral_sqrt_sum=np.zeros((img.shape[0],img.shape[1]),dtype=np.int32)
+ 
+    rows,cols=img.shape
+    for r in range(rows):
+        sum=0
+        sqrt_sum=0
+        for c in range(cols):
+            sum+=img[r][c]
+            sqrt_sum+=math.sqrt(img[r][c])
+ 
+            if r==0:
+                integral_sum[r][c]=sum
+                integral_sqrt_sum[r][c]=sqrt_sum
+            else:
+                integral_sum[r][c]=sum+integral_sum[r-1][c]
+                integral_sqrt_sum[r][c]=sqrt_sum+integral_sqrt_sum[r-1][c]
+ 
+    return integral_sum, integral_sqrt_sum
+ 
+def sauvola(img,k=0.1,kernerl=(31,31)):
+    '''
+    sauvola阈值法。
+    根据当前像素点邻域内的灰度均值与标准方差来动态计算该像素点的阈值
+    :param img:Mat--- 输入待处理图像
+    :param k:float---修正参数,一般0<k<1
+    :param kernerl:set---窗口大小
+    :return:img:Mat---阈值处理后的图像
+    '''
+    if kernerl[0]%2!=1 or kernerl[1]%2!=1:
+        raise ValueError('kernerl元组中的值必须为奇数, 请检查kernerl[0] or kernerl[1]是否为奇数!!!')
+ 
+    # 计算积分图和积分平方和图
+    integral_sum,integral_sqrt_sum=integral(img)
+    # integral_sum, integral_sqrt_sum = cv2.integral2(img)
+    # integral_sum=integral_sum[1:integral_sum.shape[0],1:integral_sum.shape[1]]
+    # integral_sqrt_sum=integral_sqrt_sum[1:integral_sqrt_sum.shape[0],1:integral_sqrt_sum.shape[1]]
+ 
+    #创建图像
+    rows,cols=img.shape
+    diff=np.zeros((rows,cols),np.float32)
+    sqrt_diff=np.zeros((rows,cols),np.float32)
+    mean=np.zeros((rows,cols),np.float32)
+    threshold=np.zeros((rows,cols),np.float32)
+    std=np.zeros((rows,cols),np.float32)
+ 
+    whalf=kernerl[0]>>1#计算领域类半径的一半
+ 
+    for row in range(rows):
+        #print('第{}行处理中...'.format(row))
+        for col in range(cols):
+            xmin=max(0,row-whalf)
+            ymin=max(0,col-whalf)
+            xmax=min(rows-1,row+whalf)
+            ymax=min(cols-1,col+whalf)
+ 
+            area=(xmax-xmin+1)*(ymax-ymin+1)
+            if area<=0:
+                sys.exit(1)
+ 
+            if xmin==0 and ymin==0:
+                diff[row,col]=integral_sum[xmax,ymax]
+                sqrt_diff[row,col]=integral_sqrt_sum[xmax,ymax]
+            elif xmin>0 and ymin==0:
+                diff[row, col] = integral_sum[xmax, ymax]-integral_sum[xmin-1,ymax]
+                sqrt_diff[row, col] = integral_sqrt_sum[xmax, ymax]-integral_sqrt_sum[xmin-1, ymax]
+            elif xmin==0 and ymin>0:
+                diff[row, col] = integral_sum[xmax, ymax] - integral_sum[xmax, ymax-1]
+                sqrt_diff[row, col] = integral_sqrt_sum[xmax, ymax] - integral_sqrt_sum[xmax, ymax-1]
+            else:
+                diagsum=integral_sum[xmax, ymax]+integral_sum[xmin-1, ymin-1]
+                idiagsum=integral_sum[xmax, ymin-1]+integral_sum[xmin-1, ymax]
+                diff[row,col]=diagsum-idiagsum
+ 
+                sqdiagsum=integral_sqrt_sum[xmax, ymax]+integral_sqrt_sum[xmin-1, ymin-1]
+                sqidiagsum=integral_sqrt_sum[xmax, ymin-1]+integral_sqrt_sum[xmin-1, ymax]
+                sqrt_diff[row,col]=sqdiagsum-sqidiagsum
+ 
+            mean[row,col]=diff[row, col]/area
+            std[row,col]=math.sqrt((sqrt_diff[row,col]-math.sqrt(diff[row,col])/area)/(area-1))
+            threshold[row,col]=mean[row,col]*(1+k*((std[row,col]/128)-1))
+ 
+            if img[row,col]<threshold[row,col]:
+                img[row,col]=0
+            else:
+                img[row,col]=255
+ 
+    return img
+
+def convert_from_cv2_to_image(img: np.ndarray) -> Image:
+    # return Image.fromarray(img)
+    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+def convert_from_image_to_cv2(img: Image) -> np.ndarray:
+    # return np.asarray(img)
+    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

nougat_api_app.py

@@ -11,7 +11,7 @@ print('GPU Device name:', torch.cuda.get_device_name(torch.cuda.current_device()
 import sys
 from functools import partial
 from http import HTTPStatus
-
+import cv2
 from fastapi import FastAPI, File, UploadFile, Request,Response, BackgroundTasks, HTTPException
 from fastapi import APIRouter, Depends
 import os 
@@ -70,8 +70,9 @@ from datetime import datetime
 from sql_app.db import Base 
 import psycopg2
 import numpy as np
+from .binarization import sauvola, convert_from_cv2_to_image, convert_from_image_to_cv2
     
-logging.basicConfig(filename='info.log', level=logging.DEBUG)
+logging.basicConfig(filename='info.log', level=logging.INFO)
 #logger = logging.getLogger()
 #logger.setLevel(logging.INFO)
 
@@ -88,6 +89,9 @@ global selected_model_name
 
 # Load the ML model
 def loadModel(checkpoint):
+    if not checkpoint.exists():
+        checkpoint = default_checkpoint_path 
+        logging.info(f"request checkpoint is not exist, using default {checkpoint_name}")       
     model = NougatModel.from_pretrained(checkpoint).to(torch.bfloat16)
     if torch.cuda.is_available():
         model.to("cuda")
@@ -242,6 +246,19 @@ async def app_middleware(request: Request, call_next):
         headers=dict(response.headers), media_type=response.media_type, background=task)
 
 ''' debug code for , not test yet    
+# Exception handlers
+def add_exception_handlers(_app: FastAPI):
+    @_app.exception_handler(ApiAuthException)
+    async def api_auth_exception_handler(request: Request, exc: ApiAuthException):
+        return await handler.api_auth_exception_handler(request, exc)
+
+    @_app.exception_handler(ApiException)
+    async def api_exception_handler(request: Request, exc: ApiException):
+        return await handler.api_exception_handler(request, exc)
+
+add_exception_handlers(main_app)
+add_exception_handlers(sub_app)
+
 from loguru import logger
 from starlette.routing import Match
 
@@ -401,6 +418,10 @@ def predict_image(model_name, images, batchsize=1, markdown=True, out_path=""):
         logging.info("we are under image to mmd convertiong")
         #sample = Image.open(images.name).convert('RGB')
         sample = images.convert('RGB')
+        gray_image = cv2.cvtColor(images, cv2.COLOR_BGR2GRAY)
+        sauvola_img = sauvola(gray_image)
+        #convert back to RGB format
+        sample = cv2.cvtColor(sauvola_img,cv2.COLOR_GRAY2RGB)
         im_new = resize_with_padding(sample, (672,896))
         img_tensor = prepare(im_new,random_padding=False)
         img_tensor = img_tensor.unsqueeze(0)
@@ -422,7 +443,7 @@ def predict_image(model_name, images, batchsize=1, markdown=True, out_path=""):
           out = out.replace(r"\(", "$").replace(r'\)', '$').replace(r'\[', '$$').replace(r'\]', '$$')
           f.write(out)
     else:
-        logging.debug(out, "\n\n")
+        logging.debug(f"the out is {out}")
 
     return model_output, [out_path]
 
@@ -523,6 +544,7 @@ def read_users(skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
 async def predict(
     request: Request, 
     selectedModel: str = Form(...), 
+    binarization: str = Form(...), 
     file: UploadFile = File(...), 
     start: int = None, stop: int = None,
 ):
@@ -549,6 +571,7 @@ async def predict(
     #parsed_url = urlparse(request)        
     #model_name = parse_qs(parsed_url.query)['model'][0]
     model_name = selectedModel
+    isBinarized = bool(binarization)
     if model_name == None:
         model = nougatModel
     else:
@@ -573,7 +596,7 @@ async def predict(
         try:
             with open(dest, 'wb') as f:
                 logging.info(f"save uploading files to {dest}")
-                imgbin = await file.read()
+                imgbin = await file.read()             
                 f.write(imgbin)
                 md5 = hashlib.md5(imgbin).hexdigest()  
                 finger_printer = md5      
@@ -584,8 +607,16 @@ async def predict(
             f.close()    
         #logging.info(f"input image type is {type(imgbin)}")
         #logging.info(f"input image type is {type(f)}")
-        
-        img = Image.open(io.BytesIO(imgbin)) 
+
+        if not isBinarized:
+            #binarize image
+            image = cv2.imread(dest)
+            gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            sauvola_img = sauvola(gray_image)
+            #convert back to RGB format
+            img = cv2.cvtColor(sauvola_img,cv2.COLOR_GRAY2RGB)       
+        else:
+            img = Image.open(io.BytesIO(imgbin)) 
         logging.info(f"uploading Image Type: {type(img)}") 
         if img.format != "PNG":
             dest_filename, dest_ext = os.path.splitext(dest)           
@@ -594,6 +625,15 @@ async def predict(
             #    f.seek(0)
             #    img = Image.open(f)
             img = convertImageFormat(img)
+            #convert to cv2 format
+            cv2_image = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+            #convert to gray format
+            gray_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
+            #binarize with sauvola algorithm 
+            sauvola_img = sauvola(gray_image)
+            #convert back to RGB format
+            img = cv2.cvtColor(sauvola_img,cv2.COLOR_GRAY2RGB)   
+            img = convert_from_cv2_to_image(img)          
             img.save( dest_filename + ".png", "PNG")  
         model_output,_ = predict_image(model_name, images=img) 
         logging.debug(f"predict output as: {model_output}")
@@ -847,7 +887,7 @@ def get_pdf(pdf_link):
 
     if response.status_code == 200:
         # Save the PDF content to a local file
-        with open(unique_filename, 'wb',encoding="utf-8") as pdf_file:
+        with open(unique_filename, 'wb') as pdf_file:
             pdf_file.write(response.content)
         logging.info("PDF downloaded successfully.")
     else:
@@ -1134,5 +1174,5 @@ if __name__ == "__main__":
     #            ssl_keyfile='/workspace/nougat-latex/lzs.chrdw.ml.key',
     #            ssl_certfile='/workspace/nougat-latex/fullchain.cer')
     
-    uvicorn.run("__main__:app", host="0.0.0.0", port=8503,log_level="debug", workers=1)   
+    uvicorn.run("__main__:app", host="0.0.0.0", port=8866,log_level="debug", workers=1)   
     #demo.launch(debug=True,share=True, server_name="0.0.0.0",server_port=8866) 

predict.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:38cbc49e9d28eb3db48cff4e23ee23122912778874db53e8891013c4c4b60744
-size 36292
+oid sha256:f7512536f5c844b04801b605568df461d30cfaa3151d4dac30878824dbb698aa
+size 4819007