CheXNet/deployment/src/streamlit_client.py

import cv2
import numpy as np
import os
import streamlit as st

from omegaconf import OmegaConf
from PIL import Image

__import__('sys').path.append('.')
from deployment.src.client import predict_api, predict_sagemaker


@st.cache
def load_image(image_file):
    pil_image = Image.open(image_file)
    image = np.array(pil_image)
    image = cv2.resize(image, (224, 224))
    image = np.float32(image / np.float32(255))
    if len(image.shape) == 2:
        image = image[..., np.newaxis].repeat(3, axis=-1)
    return image

@st.cache
def predict(image_tensor: np.ndarray, mode: str, endpoint: str):
    if mode == 'API Endpoint':
        return predict_api(image_tensor[np.newaxis, ...], api_endpoint=endpoint)[0]
    else:
        return predict_sagemaker(image_tensor[np.newaxis, ...], sagemaker_endpoint=endpoint)[0]


DEPLOYMENT_CONST_PATH = os.path.join('deployment', 'src', 'const.yaml')
GENERAL_CONST_PATH = os.path.join('src', 'const.yaml')
deployment_const = OmegaConf.load(os.path.join(os.getcwd(), DEPLOYMENT_CONST_PATH))
general_const = OmegaConf.load(os.path.join(os.getcwd(), GENERAL_CONST_PATH))

st.title('CheXNet Inference')
mode = st.selectbox('Method', ['API Endpoint', 'AWS SageMaker CLI'])

if mode == 'API Endpoint':
    api_endpoint = st.text_input('API Endpoint URL', deployment_const.API_ENDPOINT)
elif mode == 'AWS SageMaker CLI':
    sagemamer_endpoint = st.text_input('SageMaker Endpoint Name', deployment_const.SAGEMAKER_ENDPOINT)

st.markdown('---')

st.header('Select image')
image_file = st.file_uploader('Upload Image', type=['png', 'jpg', 'jpeg', 'gif', 'bmp'])

if image_file is not None:
    image = load_image(image_file)
    st.image(image)

    raw_output = predict(image, mode=mode, endpoint=api_endpoint if mode == 'API Endpoint' else sagemamer_endpoint)
    st.write("Raw model output:")
    st.write(raw_output[np.newaxis])

    mask = (raw_output >= 0.5)[:eval(general_const.N_CLASSES)]
    n_ailments = np.sum(mask)
    st.header(f"This image shows {n_ailments} ailment{'s' if n_ailments != 1 else ''}:")
    output = np.array(general_const.CLASSES)[mask]
    st.write(output)