Add Dockerfile and fix some requirement

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.8-slim-buster
+
+WORKDIR /docker_demo
+
+RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+
+RUN pip install torch==1.7.1  torchvision==0.8.2 --extra-index-url https://download.pytorch.org/whl/cpu
+
+RUN pip install Flask==2.0.3
+
+ADD . .
+
+RUN mkdir -p /root/.cache/torch/hub/checkpoints/
+
+RUN mv /docker_demo/densenet121-a639ec97.pth /root/.cache/torch/hub/checkpoints/densenet121-a639ec97.pth
+
+CMD ["python", "app.py"]

README.md

@@ -4,24 +4,37 @@ This repo contains a sample code to show how to create a Flask API server by dep
 
 If you'd like to learn how to deploy to Heroku, then check [this repo](https://github.com/avinassh/pytorch-flask-api-heroku).
 
+## Deploy with Docker
 
-## How to 
+Build image
 
-Install the dependencies:
+    docker build --network=host -t pytorch_flask:v0 .
 
-    pip install -r requirements.txt
+Run Docker
 
+    docker run --name pytorch_flask -p 5000:5000 -d pytorch_flask:v0
 
-Run the Flask server:
+send the image file in a request:
 
-    FLASK_ENV=development FLASK_APP=app.py flask run
+    curl -X POST -F file=@cat_pic.jpeg http://localhost:5000/predict
 
 
-From another tab, send the image file in a request:
+## Benchmarking with Apache Bench
 
-    curl -X POST -F file=@cat_pic.jpeg http://localhost:5000/predict
+Remove Docker:
+
+    docker stop pytorch_flask
+    docker rm pytorch_flask
+
+Run Benchmark Code:
 
+    docker run --name pytorch_flask -p 5000:5000 -d pytorch_flask:v0 python benchmark.py
+
+Benchmarking:
+
+    ab -n 100 -c 10  http://localhost:5000/predict
 
 ## License
 
 The mighty MIT license. Please check `LICENSE` for more details.
+

app.py

@@ -1,6 +1,6 @@
 import io
 import json
-
+import os
 from torchvision import models
 import torchvision.transforms as transforms
 from PIL import Image
@@ -42,4 +42,5 @@ def predict():
 
 
 if __name__ == '__main__':
-    app.run()
+    port = int(os.environ.get('PORT', 5000))
+    app.run(debug=True, host='0.0.0.0', port=port)

benchmark.py

@@ -0,0 +1,49 @@
+import io
+import json
+import os
+from torchvision import models
+import torchvision.transforms as transforms
+from PIL import Image
+from flask import Flask, jsonify, request
+
+
+app = Flask(__name__)
+imagenet_class_index = json.load(open('imagenet_class_index.json'))
+model = models.densenet121(pretrained=True)
+model.eval()
+
+
+def transform_image(image_bytes):
+    my_transforms = transforms.Compose([transforms.Resize(255),
+                                        transforms.CenterCrop(224),
+                                        transforms.ToTensor(),
+                                        transforms.Normalize(
+                                            [0.485, 0.456, 0.406],
+                                            [0.229, 0.224, 0.225])])
+    image = Image.open(io.BytesIO(image_bytes))
+    return my_transforms(image).unsqueeze(0)
+
+
+def get_prediction(image_bytes):
+    tensor = transform_image(image_bytes=image_bytes)
+    outputs = model.forward(tensor)
+    _, y_hat = outputs.max(1)
+    predicted_idx = str(y_hat.item())
+    return imagenet_class_index[predicted_idx]
+
+
+@app.route('/predict')
+def predict():
+    #with open('/docker_demo/cat_pic.jpeg', 'rb') as fd:
+    byteImgIO = io.BytesIO()
+    byteImg = Image.open("/docker_demo/cat_pic.jpeg")
+    byteImg.save(byteImgIO, "jpeg")
+    byteImgIO.seek(0)
+    img_bytes = byteImgIO.read()
+    class_id, class_name = get_prediction(image_bytes=img_bytes)
+    return jsonify({'class_id': class_id, 'class_name': class_name})
+
+
+if __name__ == '__main__':
+    port = int(os.environ.get('PORT', 5000))
+    app.run(debug=True, host='0.0.0.0', port=port)

requirements.txt

@@ -1,2 +1,3 @@
-Flask==1.0.3
-torchvision==0.3.0
+Flask==2.0.3
+torch==1.7.1 
+torchvision==0.8.2