import pre_reqs
import cv2
import numpy as np
import streamlit as st
from mmcls.apis import init_model
from mmcls.apis import inference_model_topk as inference_cls_model
from mmdet.registry import VISUALIZERS
# from mmcls.utils import register_all_modules as register_all_modules_cls
from mmdet.apis import init_detector, inference_detector
from mmdet.utils import register_all_modules as register_all_modules_det
import pandas as pd
from PIL import Image

st.set_page_config(page_title="🔥 An Image Demo", page_icon="🔥", layout='wide')
st.markdown("# 🔥 An Image Demo")
st.write(
    ":dog: Try uploading an image to get the possible categories, objects."
)
st.sidebar.header("An Image Demo")
my_upload = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
col1, col2, col3 = st.columns(3)
model_option = st.radio(
    "What\'s your inference model",
    ('cls', 'det'))

parent_folder = './'
topk = st.slider('Return top-k predictions', 1, 10, 3)

@st.cache_resource
def _init_model(model_option):
    if model_option == 'cls':
    # init model
        model = init_model(parent_folder + 'configs/resnet/resnet50_8xb32_in1k.py',
                       'https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth')
        visualizer = None
    elif model_option == 'det':
        # register_all_modules_det()
        model = init_detector(parent_folder + 'configs/rtmdet/rtmdet-ins_s_8xb32-300e_coco.py',
                          'https://download.openmmlab.com/mmdetection/v3.0/rtmdet/rtmdet-ins_s_8xb32-300e_coco/rtmdet-ins_s_8xb32-300e_coco_20221121_212604-fdc5d7ec.pth', device='cpu')
        visualizer = VISUALIZERS.build(model.cfg.visualizer)
        visualizer.dataset_meta = model.dataset_meta
    else:
        model = None
        visualizer = None
    return model, visualizer

@st.cache_data
def _get_image(my_upload=my_upload):
    if my_upload is not None:
        img_file = my_upload
    else:
        img_file = parent_folder + "images/zebra.jpg"
    return Image.open(img_file).convert('RGB')

# @st.cache_resource
def _inference_model(img, model, visualizer, model_option):
    img = np.array(img)
    if model_option == 'cls':
        return_results = inference_cls_model(model, img, 10)
        vis_img = img
    elif model_option == 'det':
        vis_img = img.copy()
        results = inference_detector(model, img)
        # import pdb
        # pdb.set_trace()
        b, h, w = results.pred_instances.masks.shape
        vis_img = cv2.resize(vis_img, (w, h))
        visualizer.add_datasample(
            name='result',
            image=vis_img,
            data_sample=results,
            draw_gt=False,
            show=False)
        vis_img = visualizer.get_image()
        cls_names = visualizer.dataset_meta['classes']
        return_results = {'scores': results.pred_instances.scores[:10].numpy(),
                          'bboxes': results.pred_instances.bboxes[:10].numpy(),
                          'labels': [cls_names[x.item()] for x in results.pred_instances.labels[:10]]
                          }
    return return_results, vis_img


def plot_canvas(img, vis_img, results, model_option):
    col1.write("Original Image :camera:")
    col1.image(img)

    col2.write("Visualization:wrench:")
    col3.write("Metainfo:wrench:")
    if model_option == 'cls':
        col2.image(vis_img)
        df = pd.DataFrame({
            'category': results["pred_class"][:topk],
            'probability': [f"{x:.2f}" for x in results["pred_score"]][:topk]
            }, index=None)
        col3.dataframe(df)
    elif model_option == 'det':
        # vis_idx = st.slider('Show a prediction', 1, 10, 3, disabled=True)
        col2.image(vis_img)

        df = pd.DataFrame({
            'category': results["labels"][:topk],
            'probability': [f"{x:.2f}" for x in results["scores"]][:topk],
            'box': [list(map(lambda t: f"{t:.2f}", list(x))) for x in results["bboxes"][:topk]]
        }, index=None)
        col3.dataframe(df)


model, visualizer = _init_model(model_option)
img = _get_image(my_upload)
results, vis_img = _inference_model(img, model, visualizer, model_option)
plot_canvas(img, vis_img, results, model_option)