Upload OmniParser/omnitool/omniparserserver/omniparserserver.py with huggingface_hub

OmniParser/omnitool/omniparserserver/omniparserserver.py

@@ -0,0 +1,70 @@
+'''
+python -m omniparserserver --som_model_path ../../weights/icon_detect/model.pt --caption_model_name florence2 --caption_model_path ../../weights/icon_caption_florence --device cuda --BOX_TRESHOLD 0.05
+'''
+
+import sys
+import os
+import time
+import json
+import base64
+import tempfile
+import shutil
+import cv2
+import numpy as np
+from pathlib import Path
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse, FileResponse
+from pydantic import BaseModel
+import argparse
+import uvicorn
+from PIL import Image
+import io
+
+root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(root_dir)
+from util.omniparser import Omniparser
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Omniparser API')
+    parser.add_argument('--som_model_path', type=str, default='weights/icon_detect/model.pt', help='Path to the som model')
+    parser.add_argument('--caption_model_name', type=str, default=None, help='Name of the caption model')
+    parser.add_argument('--caption_model_path', type=str, default='weights/icon_caption_florence', help='Path to the caption model')
+    parser.add_argument('--device', type=str, default='cpu', help='Device to run the model')
+    parser.add_argument('--BOX_TRESHOLD', type=float, default=0.05, help='Threshold for box detection')
+    parser.add_argument('--host', type=str, default='127.0.0.1', help='Host for the API')
+    parser.add_argument('--port', type=int, default=8000, help='Port for the API')
+    parser.add_argument('--save_cropped_images', action='store_true', help='Save cropped UI images to a folder')
+    parser.add_argument('--cropped_images_dir', type=str, default='cropped_images', help='Directory to save cropped images')
+    args = parser.parse_args()
+    return args
+
+args = parse_arguments()
+config = vars(args)
+
+# Convert relative paths to absolute paths
+if not os.path.isabs(config['som_model_path']):
+    config['som_model_path'] = os.path.normpath(os.path.join(root_dir, config['som_model_path']))
+if not os.path.isabs(config['caption_model_path']):
+    config['caption_model_path'] = os.path.normpath(os.path.join(root_dir, config['caption_model_path']))
+
+app = FastAPI()
+omniparser = Omniparser(config)
+
+class ParseRequest(BaseModel):
+    base64_image: str
+
+@app.post("/parse/")
+async def parse(parse_request: ParseRequest):
+    print('start parsing...')
+    start = time.time()
+    dino_labled_img, parsed_content_list = omniparser.parse(parse_request.base64_image)
+    latency = time.time() - start
+    print('time:', latency)
+    return {"som_image_base64": dino_labled_img, "parsed_content_list": parsed_content_list, 'latency': latency}
+
+@app.get("/probe/")
+async def root():
+    return {"message": "Omniparser API ready"}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=args.host, port=args.port, reload=False)