Fixed minor typos

Dockerfile

@@ -16,4 +16,4 @@ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 COPY . /code
 
-CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,4 +1,143 @@
-import subprocess
+#!/usr/bin/env python
 
-# Start HF-MaskCut with uvicorn
-subprocess.run("uvicorn server:app --host 0.0.0.0 --port 7860", shell=True)
+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 API
+import uvicorn
+
+app = FastAPI(
+    title="AdVisual MaskCut Model",
+    description="Description of the ML Model",
+    version="0.0.1",
+    terms_of_service=None,
+    contact=None,
+    license_info=None,
+    docs_url="/",
+)
+
+# 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)

requirements.txt

@@ -9,4 +9,4 @@ tqdm==4.64.1
 pycocotools==2.0.6
 fastapi==0.94.0
 pydantic==1.8.2
-uvicorn==0.21.0
+uvicorn[standard]==0.21.0

server.py

@@ -1,136 +0,0 @@
-#!/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
-
-app = FastAPI(
-    title="AdVisual MaskCut Model",
-    description="Description of the ML Model",
-    version="0.0.1",
-    terms_of_service=None,
-    contact=None,
-    license_info=None,
-    docs_url="/",
-)
-
-# 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