Merge commit 'f2b14084fb00a1f90e124484b0bf3da4f93ca311'

yjmm10 · Aug 13, 2024 · 00502fe · 00502fe
2 parents 6a0ebbf + f2b1408
commit 00502fe
Show file tree

Hide file tree

Showing 46 changed files with 1,066 additions and 710 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -30,15 +30,16 @@ jobs:
         with:
           python-version: ${{ matrix.python }}
       - name: Install Git LFS
-        if: runner.os != 'Windows'
+        shell: pwsh
         run: |
-          curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
-          sudo apt-get install git-lfs
-          git lfs install
-      - name: Install Git LFS (Windows)
-        if: runner.os == 'Windows'
-        run: |
-          choco install git-lfs
+          if ($env:RUNNER_OS -eq "Linux") {
+            curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
+            sudo apt-get install git-lfs
+          } elseif ($env:RUNNER_OS -eq "macOS") {
+            brew install git-lfs
+          } elseif ($env:RUNNER_OS -eq "Windows") {
+            choco install git-lfs
+          }
           git lfs install
       - name: Install dependencies
         run: |
@@ -53,7 +54,7 @@ jobs:
         run: |
           git config --global credential.helper store
           echo "https://USER:[email protected]" > ~/.git-credentials
-          git clone https://huggingface.co/liferecors/Telos.git telos/models
+          git clone https://huggingface.co/liferecords/Telos.git telos/models
       - name: Run tests
         run: |
           tox -e py
@@ -68,15 +69,22 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install tox isort black
+          pip install tox isort black pylint
       - name: Run isort
-        run: isort .
+        run: isort . --check-only --diff
       - name: Run black
         run: black .
+      # - name: Run pylint
+      #   run: pylint **/*.py --exit-zero --output-format=parseable --reports=y > pylint_report.txt
       - name: Run linting
         run: |
-          tox -e isort
-          tox -e pylint
+          # tox -e isort
+          # tox -e pylint
+      # - name: Upload pylint report
+      #   uses: actions/upload-artifact@v2
+      #   with:
+      #     name: pylint-report
+      #     path: pylint_report.txt
       - name: Commit changes
         uses: stefanzweifel/git-auto-commit-action@v4
         with:

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -27,3 +27,13 @@ repos:
         entry: pylint
         language: system
         types: [ python ]
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+  - repo: https://github.com/psf/black
+    rev: 23.7.0
+    hooks:
+      - id: black
+        language_version: python3.10
diff --git a/Telos_test.py b/Telos_test.py
@@ -13,24 +13,26 @@
 # result = telos.CVModel("telos")
 # print(result)
 
+
 def test_yolo():
     import cv2
+
     model_path = "detection/yolov8n_cdla.onnx"
 
     # Initialize YOLOv8 object detector
     labels = [
-            "Header",
-            "Text",
-            "Reference",
-            "Figure caption",
-            "Figure",
-            "Table caption",
-            "Table",
-            "Title",
-            "Footer",
-            "Equation",
-            ]
-    model = YOLOv8(model_path, labels=labels,conf_thres=0.3, iou_thres=0.5)
+        "Header",
+        "Text",
+        "Reference",
+        "Figure caption",
+        "Figure",
+        "Table caption",
+        "Table",
+        "Title",
+        "Footer",
+        "Equation",
+    ]
+    model = YOLOv8(model_path, labels=labels, conf_thres=0.3, iou_thres=0.5)
 
     # img_url = "https://live.staticflickr.com/13/19041780_d6fd803de0_3k.jpg"
     # img = imread_from_url(img_url)
@@ -40,16 +42,15 @@ def test_yolo():
     model(img)
 
     # Draw detections
-    combined_img = model.draw_detections(img,mask_alpha=0.2)
+    combined_img = model.draw_detections(img, mask_alpha=0.2)
     cv2.imwrite("output1.jpg", combined_img)
 
 
-
 def test_layout():
     import cv2
 
     from telos import Layout
-   
+
     model = Layout(conf_thres=0.3, iou_thres=0.5)
     img = cv2.imread("test_img/page_p6.png")
 
@@ -59,27 +60,29 @@ def test_layout():
     # print(result_T)
 
     # Draw detections
-    combined_img = model.draw_detections(img,mask_alpha=0.2)
+    combined_img = model.draw_detections(img, mask_alpha=0.2)
     cv2.imwrite("tests/output/output-layout.jpg", combined_img)
 
+
 def test_formula():
     import cv2
 
     from telos import DetFormula
-   
+
     model = DetFormula(conf_thres=0.3, iou_thres=0.5)
     img = cv2.imread("test_img/formula_page0.jpg")
 
     # Detect Objects
-    result= model(img)
+    result = model(img)
     # print(result)
     result_T = model._telos()
     print(result_T)
 
     # Draw detections
-    combined_img = model.draw_detections(img,mask_alpha=0.2)
+    combined_img = model.draw_detections(img, mask_alpha=0.2)
     cv2.imwrite("tests/output/output-formula.jpg", combined_img)
 
+
 def test_latexocr():
     import cv2
 
@@ -102,32 +105,37 @@ def test_latexocr():
     # # Draw detections
     # combined_img = model.draw_detections(img,mask_alpha=0.2)
     # cv2.imwrite("tests/output/output-formula.jpg", combined_img)
+
+
 def test_dbnet():
     import cv2
 
     from telos import DBNet
+
     model_path = "detection/det_text.onnx"
 
     # Initialize YOLOv8 object detector
-    model = DBNet(model_path,labels=["text"])
+    model = DBNet(model_path, labels=["text"])
     img = cv2.imread("/home/zyj/project/MOP/test_img/page_p0.png")
     result = model(img)
     # print(result)
 
     # Draw detections
-    combined_img = model.draw_detections(img,mask_alpha=0.2)
+    combined_img = model.draw_detections(img, mask_alpha=0.2)
     cv2.imwrite("tests/output/output-text.jpg", combined_img)
 
+
 def test_crnnnet():
     import cv2
 
     from telos import CRNN
+
     model_path = "/home/zyj/project/MOP/telos/models/recognition/rec_text"
 
     # Initialize YOLOv8 object detector
     model = CRNN(model_path="recognition/rec_text")
     img = cv2.imread("/home/zyj/project/MOP/test_img/page_p0.png")
-    result = model([img]*3)
+    result = model([img] * 3)
     print(result)
 
     result = model(img)
@@ -136,6 +144,7 @@ def test_crnnnet():
     # combined_img = model.draw_detections(img,mask_alpha=0.2)
     # cv2.imwrite("tests/output/output-text.jpg", combined_img)
 
+
 def test_OCR():
     import cv2
 
@@ -155,6 +164,7 @@ def test_OCR():
     # combined_img = model.draw_detections(img,mask_alpha=0.2)
     # cv2.imwrite("tests/output/output-text.jpg", combined_img)
 
+
 def test_OCR_wo_det():
     import cv2
 
@@ -163,12 +173,13 @@ def test_OCR_wo_det():
     # Initialize YOLOv8 object detector
     model = OCR()
     img = cv2.imread("/nas/projects/Github/Telos/tests/test_img/test_crnnnet.png")
-    result = model(img,use_det=False)
-    
+    result = model(img, use_det=False)
+
     # TODO:可视化
     res_telos = model._telos()
     print(model._json())
 
+
 def test_reading_order():
     import cv2
 
@@ -181,7 +192,7 @@ def test_reading_order():
 
     # Detect Objects
     result = det(img)
-    bboxs,scores,class_id = result
+    bboxs, scores, class_id = result
     # print(result_T)
 
     # # Draw detections
@@ -190,7 +201,7 @@ def test_reading_order():
 
     model = ReadingOrder()
     # img = cv2.imread("/home/zyj/project/MOP/test_img/page_p0.png")
-    result = model(img,bboxs)
+    result = model(img, bboxs)
     # print(model._json())
     print(result)
     # print(res_telos)
@@ -199,6 +210,7 @@ def test_reading_order():
     # combined_img = model.draw_detections(img,mask_alpha=0.2)
     # cv2.imwrite("tests/output/output-text.jpg", combined_img)
 
+
 def test_Table():
     import cv2
 

diff --git a/api/telos_api.py b/api/telos_api.py
@@ -1,11 +1,14 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse, FileResponse
 import cv2
 import numpy as np
-from telos import YOLOv8, Layout, DetFormula, LatexOCR, DBNet, CRNN, OCR, ReadingOrder, Table_TSR
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse, JSONResponse
+
+from telos import (CRNN, OCR, DBNet, DetFormula, LatexOCR, Layout,
+                   ReadingOrder, Table_TSR, YOLOv8)
 
 app = FastAPI()
 
+
 @app.post("/yolov8")
 async def yolov8_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -14,19 +17,28 @@ async def yolov8_endpoint(file: UploadFile = File(...)):
 
     model_path = "detection/yolov8n_cdla.onnx"
     labels = [
-        "Header", "Text", "Reference", "Figure caption", "Figure",
-        "Table caption", "Table", "Title", "Footer", "Equation"
+        "Header",
+        "Text",
+        "Reference",
+        "Figure caption",
+        "Figure",
+        "Table caption",
+        "Table",
+        "Title",
+        "Footer",
+        "Equation",
     ]
     model = YOLOv8(model_path, labels=labels, conf_thres=0.3, iou_thres=0.5)
     boxes, scores, class_ids = model(img)
 
     result = {
         "boxes": boxes.tolist(),
         "scores": scores.tolist(),
-        "class_ids": class_ids.tolist()
+        "class_ids": class_ids.tolist(),
     }
     return JSONResponse(content=result)
 
+
 @app.post("/layout")
 async def layout_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -39,6 +51,7 @@ async def layout_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result, "result_T": result_T})
 
+
 @app.post("/formula_detection")
 async def formula_detection_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -51,6 +64,7 @@ async def formula_detection_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result, "result_T": result_T})
 
+
 @app.post("/latex_ocr")
 async def latex_ocr_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -62,6 +76,7 @@ async def latex_ocr_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result})
 
+
 @app.post("/dbnet")
 async def dbnet_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -74,6 +89,7 @@ async def dbnet_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result})
 
+
 @app.post("/crnn")
 async def crnn_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -85,6 +101,7 @@ async def crnn_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result})
 
+
 @app.post("/ocr")
 async def ocr_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -96,6 +113,7 @@ async def ocr_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result})
 
+
 @app.post("/reading_order")
 async def reading_order_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -110,6 +128,7 @@ async def reading_order_endpoint(file: UploadFile = File(...)):
 
     return JSONResponse(content={"result": result})
 
+
 @app.post("/table_tsr")
 async def table_tsr_endpoint(file: UploadFile = File(...)):
     contents = await file.read()
@@ -124,5 +143,8 @@ async def table_tsr_endpoint(file: UploadFile = File(...)):
 
 if __name__ == "__main__":
     import uvicorn
+
     # 定义你的模型路径和字典路径
-    uvicorn.run("telos_api:app", host="0.0.0.0", port=8000, reload=True, timeout_keep_alive=30)
+    uvicorn.run(
+        "telos_api:app", host="0.0.0.0", port=8000, reload=True, timeout_keep_alive=30
+    )