chore: add gpu dockerfile

Author: jsbroks

Dockerfile

@@ -1,24 +1,3 @@
-# FROM python:3.6
-
-# RUN apt-get update
-# RUN apt-get install -y git libsm6 libxrender1 libfontconfig1
-
-# WORKDIR /workspace
-
-# COPY ./requirements_container.txt /workspace
-# # Install python package dependices
-# RUN pip install -r requirements_container.txt
-
-# WORKDIR /workspace/
-
-
-# COPY ./dextr_pb2.py /workspace
-# COPY ./dextr_pb2_grpc.py /workspace
-# COPY ./server.py /workspace
-# COPY ./dextr.proto /workspace
-
-# ENTRYPOINT [ "python", "server.py" ]
-
 FROM python:3.8.5
 
 RUN apt-get update
@@ -36,6 +15,6 @@ RUN pip install gunicorn==20.0.4
 COPY server.py /workspace
 
 EXPOSE 8000
-
+ENV DEVICE=cpu
 WORKDIR /workspace
 CMD [ "gunicorn", "-w 6", "-b 0.0.0.0:8000", "server:app" ]

Dockerfile.gpu

@@ -0,0 +1,21 @@
+FROM pytorch/pytorch:1.7.0-cuda11.0-cudnn8-runtime
+
+
+RUN apt-get update
+RUN apt-get install 'ffmpeg'\
+    'libsm6'\ 
+    'libxext6' -y
+
+RUN apt install liblzma-dev
+
+WORKDIR /workspace
+COPY requirements_container.txt /workspace
+RUN pip install -r requirements_container.txt
+RUN python -c "from dextr.model import DextrModel; DextrModel.pascalvoc_resunet101()"
+RUN pip install gunicorn==20.0.4
+COPY server.py /workspace
+
+EXPOSE 8000
+ENV DEVICE=cuda:0
+WORKDIR /workspace
+CMD [ "gunicorn", "-w 6", "-b 0.0.0.0:8000", "server:app" ]

server.py

@@ -4,12 +4,17 @@
 from PIL import Image
 import numpy as np
 import time
+import torch
+import os
 
 from imantics import Mask
 
+device = os.getenv("DEVICE", "cpu")
+torch_device = torch.device(device)
 
 app = Flask(__name__)
 model = DextrModel.pascalvoc_resunet101()
+model.eval()
 
 
 @app.route("/", methods=["GET", "POST"])
@@ -25,8 +30,11 @@ def hello_world():
         image = Image.open(path)
         print(f"Image Size: {image.size}", flush=True)
 
+        # points come in [x,y] order; this must be flipped
+        points = points[:, ::-1]
         mask = model.predict([image], [points])[0]
-        polygons = Mask(mask).polygons().points
+        mask_bin = mask >= 0.5
+        polygons = Mask(mask_bin).polygons().points
         polygons = [polygon.tolist() for polygon in polygons if len(polygon) > 2]
         print(f"Result: {polygons}", flush=True)