Update app.py

app.py

@@ -5,12 +5,13 @@ import base64
 import json
 import uuid
 import cv2
+import io
 import numpy as np
 import gradio as gr
-import shutil
+from PIL import Image, ExifTags, UnidentifiedImageError
 from time import gmtime, strftime
 from pydantic import BaseModel
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException, Request
 from fastapi.responses import JSONResponse
 from typing import Dict
 
@@ -18,67 +19,57 @@ from engine.header import *
 
 file_path = os.path.abspath(__file__)
 root_path = os.path.dirname(file_path)
-dump_path = os.path.join(root_path, "dump2/")
 
 device_id = get_deviceid().decode('utf-8')
 print_info('\t <Hardware ID> \t\t {}'.format(device_id))
 
-def activate_sdk():
-    online_key = os.environ.get("LICENSE_KEY")
-    offline_key_path = os.path.join(root_path, "license.txt")
-    
-    dict_path = os.path.join(root_path, "engine/bin")
-
+def activate_id_live_sdk():
+    id_live_key = os.environ.get("LICENSE_KEY")
+    id_live_dict_path = os.path.join(root_path, "engine/model")
     ret = -1
-    if online_key is None:
-        print_warning("Online license key not found!")
+    if id_live_key is None:
+        print_warning("ID LIVE license key not found!")
+        return ret
     else:
-        activate_ret = set_activation(online_key.encode('utf-8')).decode('utf-8')
-        ret = json.loads(activate_ret).get("errorCode", None)
+        ret = set_activation(id_live_key.encode('utf-8'))
 
     if ret == 0:
-        print_log("Successfully online activation SDK!")
-    else:
-        print_error(f"Failed to online activation SDK, Error code {ret}\n Trying offline activation SDK...");
-        if os.path.exists(offline_key_path) is False:
-            print_warning("Offline license key file not found!")
-            print_error(f"Falied to offline activation SDK, Error code {ret}")
-            return ret
+        ret = init_sdk(id_live_dict_path.encode('utf-8'))
+        if ret == 0:
+            print_log("Successfully init ID LIVE SDK!")
         else:
-            file=open(offline_key_path,"r")
-            offline_key = file.read()
-            file.close()
-            activate_ret = set_activation(offline_key.encode('utf-8')).decode('utf-8')
-            ret = json.loads(activate_ret).get("errorCode", None)
-            if ret == 0:
-                print_log("Successfully offline activation SDK!")
-            else:
-                print_error(f"Falied to offline activation SDK, Error code {ret}")
-                return ret
-    
-    init_ret = init_sdk(dict_path.encode('utf-8')).decode('utf-8')
-    ret = json.loads(activate_ret).get("errorCode", None)
-    print_log(f"Init SDK: {ret}")
+            print_log("Failed to init ID LIVE SDK!")
+    else:
+        print_error(f"Falied to activate ID LIVE SDK, Error code {ret}")
+
     sys.stdout.flush()
     return ret
 
 
-async def save_upload_file(upload_file: UploadFile) -> str:
-    file_name = uuid.uuid4().hex[:6] + "_" + upload_file.filename
-    save_path = os.path.join(dump_path, file_name)
-    with open(save_path, "wb") as buffer:
-        shutil.copyfileobj(upload_file.file, buffer)
-    return os.path.abspath(save_path)
+def apply_exif_rotation(image):
+    try:
+        exif = image._getexif()
+        if exif is not None:
+            for orientation in ExifTags.TAGS.keys():
+                if ExifTags.TAGS[orientation] == 'Orientation':
+                    break
+
+            # Get the orientation value
+            orientation = exif.get(orientation, None)
+
+            # Apply the appropriate rotation based on the orientation
+            if orientation == 3:
+                image = image.rotate(180, expand=True)
+            elif orientation == 6:
+                image = image.rotate(270, expand=True)
+            elif orientation == 8:
+                image = image.rotate(90, expand=True)
+
+    except AttributeError:
+        print("No EXIF data found")
+
+    return image
 
-async def save_base64_file(base64_string: str) -> str:
-    file_name = uuid.uuid4().hex[:6] + ".jpg"  # or ".png" depending on your use
-    save_path = os.path.join(dump_path, file_name)
-    
-    with open(save_path, "wb") as buffer:
-        buffer.write(base64.b64decode(base64_string))
-    
-    return os.path.abspath(save_path)
-                
 app = FastAPI()
 
 class ImageBase64Request(BaseModel):
@@ -88,112 +79,49 @@ class ImageBase64Request(BaseModel):
 def read_root():
     return {"status": "API is running"}
 
-@app.post("/api/read_idcard")
-async def read_idcard(image: UploadFile = File(...), image2: UploadFile = File(None)):
-    try:
-        file_path1 = await save_upload_file(image)
-
-        file_path2 = ""
-        if image2 is not None:
-            file_path2 = await save_upload_file(image2)
-
-
-        ocrResult = ocr_id_card(file_path1.encode('utf-8'), file_path2.encode('utf-8'))
-        ocrResDict = json.loads(ocrResult)  
-        
-        os.remove(file_path1)
-        if file_path2:
-            os.remove(file_path2)
-
-        return JSONResponse(content={"status": "ok", "data": ocrResDict}, status_code=200)
-
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
-
-
-@app.post("/api/read_idcard_base64")
-async def read_idcard_base64(request_data: ImageBase64Request):
+@app.post("/process_image")
+async def process_image(image: UploadFile = File(...)):
     try:
-        file_path = await save_base64_file(request_data.image)
-
-        # Your OCR function (pass bytes-encoded file path)
-        ocr_result = ocr_id_card(file_path.encode('utf-8'), ''.encode('utf-8'))
-        ocr_res_dict = json.loads(ocr_result)
-
-        # Remove the temp file
-        os.remove(file_path)
+        contents = await image.read()
+        image_pil = apply_exif_rotation(Image.open(io.BytesIO(contents))).convert('RGB')
+    except UnidentifiedImageError:
+        return JSONResponse(status_code=400, content={"resultCode": "Error", "result": "Failed to open file"})
 
-        return JSONResponse(content={"status": "ok", "data": ocr_res_dict}, status_code=200)
+    image_np = np.asarray(image_pil)
+    result = processImage(image_np, image_np.shape[1], image_np.shape[0])
 
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
-
-@app.post("/api/read_credit")
-async def read_credit(image: UploadFile = File(...)):
-    try:
-        file_path = await save_upload_file(image)
-
-        ocrResult = ocr_credit_card(file_path.encode('utf-8'))
-        ocrResDict = json.loads(ocrResult)  
-        
-        os.remove(file_path)
-
-        return JSONResponse(content={"status": "ok", "data": ocrResDict}, status_code=200)
-
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
-
-@app.post("/api/read_credit_base64")
-async def read_credit_base64(request_data: ImageBase64Request):
-    try:
-        file_path = await save_base64_file(request_data.image)
-
-        # Your OCR function (pass bytes-encoded file path)
-        ocr_result = ocr_credit_card(file_path.encode('utf-8'))
-        ocr_res_dict = json.loads(ocr_result)
-
-        # Remove the temp file
-        os.remove(file_path)
-
-        return JSONResponse(content={"status": "ok", "data": ocr_res_dict}, status_code=200)
-
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
+    if result is None:
+        return JSONResponse(status_code=400, content={"resultCode": "Error", "result": "Failed to process image"})
+    
+    result_dict = json.loads(result.decode('utf-8'))
+    return JSONResponse(status_code=200, content={"resultCode": "Ok", "result": result_dict})
 
 
-@app.post("/api/read_barcode")
-async def read_barcode(image: UploadFile = File(...)):
+@app.post("/process_image_base64")
+async def process_image_base64(request: ImageBase64Request):
     try:
-        file_path = await save_upload_file(image)
+        image_data = base64.b64decode(request.image)
+        image = apply_exif_rotation(Image.open(io.BytesIO(image_data))).convert("RGB")
+    except (base64.binascii.Error, UnidentifiedImageError, ValueError) as e:
+        raise HTTPException(status_code=400, detail=f"Failed to parse base64: {str(e)}")
 
-        ocrResult = ocr_barcode(file_path.encode('utf-8'))
-        ocrResDict = json.loads(ocrResult)  
-        
-        os.remove(file_path)
+    image_np = np.asarray(image)
+    result = processImage(image_np, image_np.shape[1], image_np.shape[0])
 
-        return JSONResponse(content={"status": "ok", "data": ocrResDict}, status_code=200)
+    if result is None:
+        raise HTTPException(status_code=400, detail="Failed to process image")
 
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
-
-@app.post("/api/read_barcode_base64")
-async def read_barcode_base64(request_data: ImageBase64Request):
     try:
-        file_path = await save_base64_file(request_data.image)
-
-        # Your OCR function (pass bytes-encoded file path)
-        ocr_result = ocr_barcode(file_path.encode('utf-8'))
-        ocr_res_dict = json.loads(ocr_result)
-
-        # Remove the temp file
-        os.remove(file_path)
-
-        return JSONResponse(content={"status": "ok", "data": ocr_res_dict}, status_code=200)
-
+        result_dict = json.loads(result.decode("utf-8"))
     except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
+        raise HTTPException(status_code=500, detail=f"Failed to decode result: {str(e)}")
 
+    return JSONResponse(
+        content={"resultCode": "Ok", "result": result_dict},
+        status_code=200
+    )
 
+    
 if __name__ == '__main__':
     ret = activate_sdk()
     if ret != 0: