From 0cfc8d1e3bead71976a0eb71f3c794c2eba1b184 Mon Sep 17 00:00:00 2001
From: sml2h3 <sml2h3@gmail.com>
Date: Thu, 25 Jul 2024 10:48:35 +0800
Subject: [PATCH] =?UTF-8?q?=E6=95=B4=E4=BD=93=E9=A1=B9=E7=9B=AE=E9=87=8D?=
 =?UTF-8?q?=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md   | 38 +++++++++++++++++++++++++++-----------
 app/main.py | 35 +++++++++++++++++++++++------------
 2 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index b3b2960..434b951 100644
--- a/README.md
+++ b/README.md
@@ -32,23 +32,39 @@
    cd ddddocr-api
    ```
 
-2. **构建 Docker 镜像 [一键docker环境服务器购买，可一元试用](https://app.rainyun.com/apps/rcs/buy) **
-   ```bash
-   docker build -t ddddocr-api .
-   ```
+2. **启动服务**
+   
+   有三种方式可以启动应用：
+
+   a. 使用 docker启动：
+      1. 构建 Docker 镜像 [一键docker环境服务器购买，可一元试用](https://app.rainyun.com/apps/rcs/buy) 
+      2. 打包镜像
+          ```bash
+          docker build -t ddddocr-api .
+          ```
+      3. 启动镜像
+         ```bash
+         docker run -d -p 8000:8000 --name ddddocr-api-container ddddocr-api
+         ```
+
+   b. 使用 python 命令直接运行：
+      ```bash
+      python app/main.py
+      ```
+   
+   b. 使用 uvicorn（支持热重载，适合开发）：
+      ```bash
+      uvicorn app.main:app --reload
+      ```
 
-3. **启动服务**
-   ```bash
-   docker run -d -p 8000:8000 --name ddddocr-api-container ddddocr-api
-   ```
 
-4. **验证服务**
+3. **验证服务**
    ```bash
    curl http://localhost:8000/docs
    ```
    > 如果成功，您将看到 Swagger UI 文档页面。
    
-5. **停止服务**
+4. **停止服务**
 
 - 如果使用 Docker：
   ```bash
@@ -60,7 +76,7 @@
   docker-compose down
   ```
   
-6. **查看日志**
+5. **查看日志**
 
 - 如果使用 Docker：
   ```bash
diff --git a/app/main.py b/app/main.py
index abf6249..c16f808 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,23 +1,30 @@
+import uvicorn
 from fastapi import FastAPI, File, UploadFile, HTTPException, Form
-from fastapi.responses import JSONResponse
 from typing import Optional, Union
 import base64
-from .models import OCRRequest, SlideMatchRequest, DetectionRequest, APIResponse
-from .services import ocr_service
+from app.models import OCRRequest, SlideMatchRequest, DetectionRequest, APIResponse
+from app.services import ocr_service
 
 app = FastAPI()
 
+from starlette.datastructures import UploadFile as StarletteUploadFile
 
-def decode_image(image: Union[UploadFile, str, None]) -> bytes:
-    if isinstance(image, UploadFile):
-        return image.file.read()
+
+async def decode_image(image: Union[UploadFile, StarletteUploadFile, str, None]) -> bytes:
+    if image is None:
+        raise HTTPException(status_code=400, detail="No image provided")
+
+    if isinstance(image, (UploadFile, StarletteUploadFile)):
+        return await image.read()
     elif isinstance(image, str):
         try:
+            # 检查是否是 base64 编码的图片
+            if image.startswith(('data:image/', 'data:application/')):
+                # 移除 MIME 类型前缀
+                image = image.split(',')[1]
             return base64.b64decode(image)
         except:
             raise HTTPException(status_code=400, detail="Invalid base64 string")
-    elif image is None:
-        raise HTTPException(status_code=400, detail="No image provided")
     else:
         raise HTTPException(status_code=400, detail="Invalid image input")
 
@@ -34,7 +41,7 @@ async def ocr_endpoint(
         if file is None and image is None:
             return APIResponse(code=400, message="Either file or image must be provided")
 
-        image_bytes = decode_image(file or image)
+        image_bytes = await decode_image(file or image)
         result = ocr_service.ocr_classification(image_bytes, probability, charsets, png_fix)
         return APIResponse(code=200, message="Success", data=result)
     except Exception as e:
@@ -53,8 +60,8 @@ async def slide_match_endpoint(
         if (target_file is None and target is None) or (background_file is None and background is None):
             return APIResponse(code=400, message="Both target and background must be provided")
 
-        target_bytes = decode_image(target_file or target)
-        background_bytes = decode_image(background_file or background)
+        target_bytes = await decode_image(target_file or target)
+        background_bytes = await decode_image(background_file or background)
         result = ocr_service.slide_match(target_bytes, background_bytes, simple_target)
         return APIResponse(code=200, message="Success", data=result)
     except Exception as e:
@@ -70,8 +77,12 @@ async def detection_endpoint(
         if file is None and image is None:
             return APIResponse(code=400, message="Either file or image must be provided")
 
-        image_bytes = decode_image(file or image)
+        image_bytes = await decode_image(file or image)
         bboxes = ocr_service.detection(image_bytes)
         return APIResponse(code=200, message="Success", data=bboxes)
     except Exception as e:
         return APIResponse(code=500, message=str(e))
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)