Added subprocess running uvicorn

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "CutLER"]
 	path = CutLER
-	url = https://github.com/facebookresearch/CutLER
+	url = https://github.com/Ad-Visual/CutLER

app.py

@@ -33,6 +33,7 @@ DESCRIPTION = 'This is an unofficial demo for https://github.com/facebookresearc
 
 paths = sorted(pathlib.Path('CutLER/maskcut/imgs').glob('*.jpg'))
 demo = gr.Interface(fn=run,
+                    enable_queue=True,
                     inputs=[
                         gr.Image(label='Input image', type='filepath'),
                         gr.Slider(

config.py

@@ -0,0 +1,62 @@
+
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import os
+import torch
+
+
+# Config that serves all environment
+GLOBAL_CONFIG = {
+    "MODEL_PATH": "../model/model.pt",
+    "SCALAR_PATH": "../model/scaler.joblib",
+    "USE_CUDE_IF_AVAILABLE": True,
+    "ROUND_DIGIT": 6
+}
+
+# Environment specific config, or overwrite of GLOBAL_CONFIG
+ENV_CONFIG = {
+    "development": {
+        "DEBUG": True
+    },
+
+    "staging": {
+        "DEBUG": True
+    },
+
+    "production": {
+        "DEBUG": False,
+        "ROUND_DIGIT": 3
+    }
+}
+
+
+def get_config() -> dict:
+    """
+    Get config based on running environment
+    :return: dict of config
+    """
+
+    # Determine running environment
+    ENV = os.environ['PYTHON_ENV'] if 'PYTHON_ENV' in os.environ else 'development'
+    ENV = ENV or 'development'
+
+    # raise error if environment is not expected
+    if ENV not in ENV_CONFIG:
+        raise EnvironmentError(f'Config for envirnoment {ENV} not found')
+
+    config = GLOBAL_CONFIG.copy()
+    config.update(ENV_CONFIG[ENV])
+
+    config['ENV'] = ENV
+    config['DEVICE'] = 'cuda' if torch.cuda.is_available() and config['USE_CUDE_IF_AVAILABLE'] else 'cpu'
+
+    return config
+
+# load config for import
+CONFIG = get_config()
+
+if __name__ == '__main__':
+    # for debugging
+    import json
+    print(json.dumps(CONFIG, indent=4))

main.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.logger import logger
+
+from model import Model
+import base64
+from io import BytesIO
+from pydantic import BaseModel
+from config import CONFIG
+
+from predict import predict
+
+# About
+import torch
+import os
+import sys
+
+# Server Framework
+import uvicorn
+
+app = FastAPI(
+    title="ML Model",
+    description="Description of the ML Model",
+    version="0.0.1",
+    terms_of_service=None,
+    contact=None,
+    license_info=None
+)
+
+# Allow CORS for local debugging
+app.add_middleware(CORSMiddleware, allow_origins=["*"])
+
+@app.on_event("startup")
+async def startup_event():
+    """
+    Initialize FastAPI and add variables
+    """
+
+    logger.info('Running envirnoment: {}'.format(CONFIG['ENV']))
+    logger.info('PyTorch using device: {}'.format(CONFIG['DEVICE']))
+
+    # Initialize the pytorch model
+    model = Model()
+
+    # add model and other preprocess tools too app state
+    app.package = {
+        "model": model
+    }
+
+@app.get("/ping")
+def ping():
+    return {"ok": True, "message": "Pong"}
+
+@app.get("/about")
+def show_about():
+    """
+    Get deployment information, for debugging
+    """
+
+    logger.info('API /about called')
+
+    def bash(command):
+        output = os.popen(command).read()
+        return output
+
+    return {
+        "sys.version": sys.version,
+        "torch.__version__": torch.__version__,
+        "torch.cuda.is_available()": torch.cuda.is_available(),
+        "torch.version.cuda": torch.version.cuda,
+        "torch.backends.cudnn.version()": torch.backends.cudnn.version(),
+        "torch.backends.cudnn.enabled": torch.backends.cudnn.enabled,
+        "nvidia-smi": bash('nvidia-smi')
+    }
+
+class ImageBody(BaseModel):
+    image: str
+    threshold: float = 0.15
+    num_objects: int = 1
+
+@app.post("/predict")
+async def do_predict(body: ImageBody):
+    """
+    Perform prediction on input data
+    """
+
+    logger.info('API predict called')
+
+    image: str = body.image
+    threshold: float = body.threshold
+    num_objects: int = body.num_objects
+
+    # Run the algorithm
+    result = predict(app.package, image, threshold, num_objects)
+
+    # Convert the result to base64 and send the json back
+    buffered = BytesIO()
+    result.save(buffered, format="JPEG")
+    img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    return {"ok": True, "status": "FINISHED", "result": img_str}
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    while True:
+        try: 
+          data = await websocket.receive_json()
+          image: str = data.get('image')
+          threshold: float = data.get('threshold') or 0.15
+          num_objects: int = data.get('num_objects') or 1
+
+          await websocket.send_json({"ok": True, "status": "STARTED"})
+
+          if image == None:
+              await websocket.send_json({
+                  "ok": False,
+                  "status": "ERROR",
+                  "message": "No image provided"
+              })
+              break
+          
+          # Run the algorithm
+          result = predict(app.package, image, threshold, num_objects)
+
+          # Convert the result to base64 and send the json back
+          buffered = BytesIO()
+          result.save(buffered, format="JPEG")
+          img_str = 'data:image/jpeg;base64,' + base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+          await websocket.send_json({"ok": True, "status": "FINISHED", "result": img_str})
+
+          await websocket.close()
+        except WebSocketDisconnect:
+          break 
+
+if __name__ == '__main__':
+    # server api
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)

model.py

@@ -1,20 +1,20 @@
 # This file is adapted from https://github.com/facebookresearch/CutLER/blob/077938c626341723050a1971107af552a6ca6697/maskcut/demo.py
 # The original license file is the file named LICENSE.CutLER in this repo.
 
+import os
 import sys
+sys.path.append('./CutLER/')
+sys.path.append('./CutLER/maskcut/')
 
 import numpy as np
 import PIL.Image as Image
 import torch
 from scipy import ndimage
 
-sys.path.append('CutLER/maskcut/')
-sys.path.append('CutLER/')
-import dino
-from colormap import random_color
-from crf import densecrf
-from maskcut import maskcut
-from third_party.TokenCut.unsupervised_saliency_detection import metric
+from CutLER.maskcut.dino import ViTFeat  # model
+from CutLER.maskcut.crf import densecrf
+from CutLER.maskcut.maskcut import maskcut
+from CutLER.third_party.TokenCut.unsupervised_saliency_detection import metric
 
 
 class Model:
@@ -33,21 +33,21 @@ class Model:
         feat_dim = 768
 
         # extract patch features with a pretrained DINO model
-        backbone = dino.ViTFeat(url, feat_dim, vit_arch, vit_feat, patch_size)
+        backbone = ViTFeat(url, feat_dim, vit_arch, vit_feat, patch_size)
         backbone.eval()
         backbone.to(self.device)
         return backbone
 
-    def __call__(self, img_path, tau, n, fixed_size=480):
+    def __call__(self, image, tau, n, fixed_size=480):
         # get pseudo-masks with MaskCut
-        bipartitions, _, I_new = maskcut(img_path,
+        bipartitions, _, I_new = maskcut(image,
                                          self.backbone,
                                          self.backbone.patch_size,
                                          tau,
                                          N=n,
                                          fixed_size=fixed_size,
                                          cpu=self.device.type == 'cpu')
-        I = Image.open(img_path).convert('RGB')
+        I = image.convert('RGB')
         width, height = I.size
         pseudo_mask_list = []
         for idx, bipartition in enumerate(bipartitions):

predict.py

@@ -0,0 +1,27 @@
+import base64
+from io import BytesIO
+from PIL import Image
+import numpy as np
+from model import Model
+
+def predict(package, image_base64: str, threshold: float, num_objects: int): 
+    # Decode the image from base64 to PIL.Image
+    # We use BytesIO to convert the base64 to bytes
+    base64_split = image_base64.split(',')[1]
+    buf = BytesIO(base64.b64decode(base64_split))
+
+    image = Image.open(buf)
+
+    # Get the image path from tmp_image
+    canvas = Image.new('RGB', image.size, (0, 0, 0))
+
+    # We copy the image that and fill it with black, to get the dimensions
+    rgb = np.array(canvas)
+    model = package.get('model')
+    masks = model(image, threshold, num_objects)
+
+    for mask in masks:  
+      fg = mask > 0.5
+      rgb[fg] = 255
+
+    return Image.fromarray(rgb)

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/lucasb-eyer/pydensecrf@0d53acb
+git+https://github.com/lucasb-eyer/pydensecrf
 gradio==3.16.2
 numpy==1.23.5
 opencv-python==4.6.0.66
@@ -7,3 +7,6 @@ scikit-image==0.19.2
 torch==1.13.1
 torchvision==0.14.1
 tqdm==4.64.1
+fastapi==0.94.0
+pydantic==1.8.2
+uvicorn==0.21.0

start.py

@@ -0,0 +1,5 @@
+import subprocess
+
+# Start HF-MaskCut with uvicorn
+
+subprocess.run(["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"])