#!/usr/bin/env python
# coding: utf-8

# # Fastai 2 medical imaging
# 
# **BentoML makes moving trained ML models to production easy:**
# 
# * Package models trained with **any ML framework** and reproduce them for model serving in production
# * **Deploy anywhere** for online API serving or offline batch serving
# * High-Performance API model server with *adaptive micro-batching* support
# * Central hub for managing models and deployment process via Web UI and APIs
# * Modular and flexible design making it *adaptable to your infrastrcuture*
# 
# BentoML is a framework for serving, managing, and deploying machine learning models. It is aiming to bridge the gap between Data Science and DevOps, and enable teams to deliver prediction services in a fast, repeatable, and scalable way.
# 
# Before reading this example project, be sure to check out the [Getting started guide](https://github.com/bentoml/BentoML/blob/master/guides/quick-start/bentoml-quick-start-guide.ipynb) to learn about the basic concepts in BentoML.
# 
# This notebook demonstrates using BentoML to deploy Fastai2 medical image classification model as API server.
# 
# ![Impression](https://www.google-analytics.com/collect?v=1&tid=UA-112879361-3&cid=555&t=event&ec=fast-ai&ea=fast-ai2-medical-imaging&dt=fast-ai2-medical-imaging)
# 

# In[1]:


get_ipython().run_line_magic('reload_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
get_ipython().run_line_magic('matplotlib', 'inline')


# In[23]:


get_ipython().system('pip install -U -q fastai>=2.0.13 pydicom>=1.4.0 nbdev>=1.0.0 kornia>=0.4.0 fastcore>=1.0.13')


# In[2]:


from fastai.data.block import DataBlock, CategoryBlock
from fastai.data.external import untar_data, URLs
from fastai.data.transforms import RandomSplitter
from fastai.metrics import accuracy
from fastai.vision.data import ImageBlock
from fastai.medical.imaging import get_dicom_files, dcmread, PILDicom
from fastai.vision.augment import aug_transforms
from fastai.vision.learner import cnn_learner
from fastai.vision.models import resnet34
# from fastai.basics import *
# from fastai.callback.all import *
# from fastai.vision.all import *
# from fastai.medical.imaging import *


import pydicom

import pandas as pd

import torch
import bentoml.utils.cloudpickle as cloudpickle


# In[3]:


pneumothorax_source = untar_data(URLs.SIIM_SMALL)


# In[4]:


items = get_dicom_files(pneumothorax_source/f"train/")
trn, val = RandomSplitter()(items)


# In[5]:


patient = 3
xray_sample = dcmread(items[patient])
xray_sample.show()


# In[6]:


df = pd.read_csv(pneumothorax_source/f"labels.csv")
df.head()


# In[9]:


get_ipython().run_cell_magic('writefile', 'datablock_utils.py', '\nfrom fastai.data.external import untar_data, URLs\n\npneumothorax_source = untar_data(URLs.SIIM_SMALL)\n\ndef get_x(x):\n    return pneumothorax_source/f"{x[0]}"\n\ndef get_y(x):\n    return x[1]\n')


# In[10]:


from datablock_utils import get_x, get_y

pneumothorax = DataBlock(blocks=(ImageBlock(cls=PILDicom), CategoryBlock),
                   #get_x=lambda x:pneumothorax_source/f"{x[0]}",
                   #get_y=lambda x:x[1],
                   get_x=get_x,
                   get_y=get_y,
                   batch_tfms=aug_transforms(size=224))


# In[11]:


dls = pneumothorax.dataloaders(df.values)
dls.show_batch(max_n=16)


# In[12]:


learn = cnn_learner(dls, resnet34, metrics=accuracy)
learn.fit_one_cycle(1)


# In[13]:


learn.predict(pneumothorax_source/f"train/Pneumothorax/000004.dcm")


# In[14]:


get_ipython().run_cell_magic('writefile', 'medical_image_service.py', "from bentoml.frameworks.fastai import FastaiModelArtifact\nfrom bentoml.adapters import FileInput\nfrom fastcore.utils import tuplify, detuplify\n\nimport bentoml\nimport datablock_utils\n\n@bentoml.artifacts([FastaiModelArtifact('learner')])\n@bentoml.env(infer_pip_packages=True)\nclass FastaiMedicalImagingService(bentoml.BentoService):\n\n    @bentoml.api(input=FileInput(), batch=True)\n    def predict(self, files):\n        files = [i.read() for i in files]\n        dl = self.artifacts.learner.dls.test_dl(files, rm_type_tfms=None, num_workers=0)\n        inp, preds, _, dec_preds = self.artifacts.learner.get_preds(dl=dl, with_input=True, with_decoded=True)\n        i = getattr(self.artifacts.learner.dls, 'n_inp', -1)\n        inp = (inp,)\n        dec_list = self.artifacts.learner.dls.decode_batch(inp + tuplify(dec_preds))\n        res = []\n        for dec in dec_list:\n            dec_inp, dec_targ = map(detuplify, [dec[:i], dec[i:]])\n            res.append(dec_targ)\n        return res\n")


# In[16]:


import bentoml

# from fastcore.utils import remove_patches_path
from medical_image_service import FastaiMedicalImagingService

learn.metrics = []

svc = FastaiMedicalImagingService()
svc.pack('learner', learn)

# Fastcore modify Pathlib module from python, need to remove the patch in order to save
# with remove_patches_path():
saved_path = svc.save()


# In[16]:


bento_tag = f'{svc.name}:{svc.version}'


# ## REST API Model Serving
# 
# To start a REST API model server with the BentoService saved above, use the bentoml serve command:

# In[ ]:


get_ipython().system('bentoml serve FastaiMedicalImagingService:latest')


# If you are running this notebook from Google Colab, you can start the dev server with --run-with-ngrok option, to gain acccess to the API endpoint via a public endpoint managed by [ngrok](https://ngrok.com):

# In[ ]:


get_ipython().system('bentoml serve FastaiMedicalImagingService:atest --run-with-ngrok')


# ### Send prediction request to the REST API server
# 
# Go visit http://127.0.0.1:5000/ from your browser, click /predict -> Try it out -> Choose File -> Execute to sumbit an image from your computer
# 
# or Using `curl` request:
# 
# ```
# curl -X POST "http://localhost:5000/predict" \
#   -F image=@000004.dcm
# ```

# ### Containerize model server with Docker
# 
# 
# One common way of distributing this model API server for production deployment, is via Docker containers. And BentoML provides a convenient way to do that.
# 
# Note that docker is **not available in Google Colab**. You will need to download and run this notebook locally to try out this containerization with docker feature.
# 
# If you already have docker configured, simply run the follow command to product a docker container serving the IrisClassifier prediction service created above:

# In[23]:


get_ipython().system('bentoml containerize FastaiMedicalImagingService:latest')


# In[24]:


get_ipython().system('docker run --rm -p 5000:5000 fastaimedicalimagingservice:20200922205913_3AE93C')


# ### Launch inference job from CLI
# 
# 
# BentoML cli supports loading and running a packaged model from CLI. With the DataframeInput adapter, the CLI command supports reading input Dataframe data from CLI argument or local csv or json files:

# In[ ]:


get_ipython().system('bentoml run FastaiMedicalImagingService:latest predict      --input-file /Users/bozhaoyu/.fastai/data/siim_small/train/Pneumothorax/000004.dcm')


# ### Load saved BentoService
# 
# bentoml.load is the API for loading a BentoML packaged model in python:

# In[ ]:


from bentoml import load


loaded_svc = load(saved_path)
print(loaded_svc.predict(pneumothorax_source/f"train/Pneumothorax/000004.dcm"))


# # Deployment Options
# 
# If you are at a small team with limited engineering or DevOps resources, try out automated deployment with BentoML CLI, currently supporting AWS Lambda, AWS SageMaker, and Azure Functions:
# - [AWS Lambda Deployment Guide](https://docs.bentoml.org/en/latest/deployment/aws_lambda.html)
# - [AWS SageMaker Deployment Guide](https://docs.bentoml.org/en/latest/deployment/aws_sagemaker.html)
# - [Azure Functions Deployment Guide](https://docs.bentoml.org/en/latest/deployment/azure_functions.html)
# 
# If the cloud platform you are working with is not on the list above, try out these step-by-step guide on manually deploying BentoML packaged model to cloud platforms:
# - [AWS ECS Deployment](https://docs.bentoml.org/en/latest/deployment/aws_ecs.html)
# - [Google Cloud Run Deployment](https://docs.bentoml.org/en/latest/deployment/google_cloud_run.html)
# - [Azure container instance Deployment](https://docs.bentoml.org/en/latest/deployment/azure_container_instance.html)
# - [Heroku Deployment](https://docs.bentoml.org/en/latest/deployment/heroku.html)
# 
# Lastly, if you have a DevOps or ML Engineering team who's operating a Kubernetes or OpenShift cluster, use the following guides as references for implementating your deployment strategy:
# - [Kubernetes Deployment](https://docs.bentoml.org/en/latest/deployment/kubernetes.html)
# - [Knative Deployment](https://docs.bentoml.org/en/latest/deployment/knative.html)
# - [Kubeflow Deployment](https://docs.bentoml.org/en/latest/deployment/kubeflow.html)
# - [KFServing Deployment](https://docs.bentoml.org/en/latest/deployment/kfserving.html)
# - [Clipper.ai Deployment Guide](https://docs.bentoml.org/en/latest/deployment/clipper.html)