Pytorch-Docker-Deployment

app.py

from flask import Flask, request, jsonify
import numpy as np
import cv2 as cv
import traceback
import json

from predict_l import get_prediction_l
from predict_x import get_prediction_x

app = Flask(__name__)

@app.route('/')
def index():
    return jsonify('Serving Yolov5 Pytorch models')

@app.route('/prediction_l', methods = ['POST'])
def prediction_l():

    try:
        f1 = request.files['image']
        f2 = np.fromstring(f1.read(), np.uint8)
        f3 = cv.imdecode(f2, cv.IMREAD_COLOR)

        res = get_prediction_l(f3)

    except Exception as e:
        print(str(traceback.format_exc()))

    return json.dumps(res)

@app.route('/prediction_x', methods = ['POST'])
def prediction_x():

    try:
        f1 = request.files['image']
        f2 = np.fromstring(f1.read(), np.uint8)
        f3 = cv.imdecode(f2, cv.IMREAD_COLOR)

        res = get_prediction_x(f3)

    except Exception as e:
        print(str(traceback.format_exc()))

    return json.dumps(res)

if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=True)

docker-compose.yml

version: "3.3"             

services:

  # Flask App
  app1-yolov5:
    build: ./app1
    image: "app1-yolov5"
    #restart: always
    container_name: app1-yolov5
    ports:
      - "5001:5000"
    deploy:
      replicas: 1
      resources:
        limits:
          memory: 5G 
        reservations:
          memory: 3G

Dockerfile

FROM ubuntu:18.04

RUN apt-get update \
    && apt-get install -y apt-utils \
    python3.8 \
    python3-pip \
    libgl1-mesa-glx \
    nano

WORKDIR /app
COPY . /app

RUN pip3 install --upgrade pip
RUN pip3 install numpy>=1.18.5
RUN pip3 install Cython
RUN pip3 install pycocotools>=2.0 
RUN pip3 install -r requirements.txt

ENTRYPOINT ["python3"]
CMD ["app.py"]

predict_l.py

import torch
import os

os.environ["CUDA_VISIBLE_DEVICES"]=""

model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='model/yolov5l.pt')
model.eval()

def get_prediction_l(image):

    # Inference
    with torch.no_grad():
        results = model(image)

    # Results
    results.print()

    # Data
    pred_res = results.xyxy[0].tolist()

    return pred_res

predict_x.py

import torch
import os

os.environ["CUDA_VISIBLE_DEVICES"]=""

model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='model/yolov5x.pt')
model.eval()

def get_prediction_x(image):

    # Inference
    with torch.no_grad():
        results = model(image)

    # Results
    results.print()

    # Data
    pred_res = results.xyxy[0].tolist()

    return pred_res

requirements.txt

# base ----------------------------------------
matplotlib>=3.2.2
opencv-python>=4.1.2
Pillow
PyYAML>=5.3.1
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0

# logging -------------------------------------
tensorboard>=2.4.1
# wandb

# plotting ------------------------------------
seaborn>=0.11.0
pandas

# export --------------------------------------
# coremltools>=4.1
# onnx>=1.8.1
# scikit-learn==0.19.2  # for coreml quantization

# extras --------------------------------------
thop  # FLOPS computation
Flask