import os
os.system('pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')
os.system("git clone https://github.com/microsoft/unilm.git")

import sys
sys.path.append("unilm")

import cv2

from unilm.dit.object_detection.ditod import add_vit_config

import torch
import numpy as np

from detectron2.config import CfgNode as CN
from detectron2.config import get_cfg
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor

import gradio as gr


# Step 1: instantiate config
cfg = get_cfg()
add_vit_config(cfg)
cfg.merge_from_file("cascade_dit_base.yml")

# Step 2: add model weights URL to config
cfg.MODEL.WEIGHTS = "https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_cascade.pth"

# Step 3: set device
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Step 4: define model
predictor = DefaultPredictor(cfg)

def get_bytes_shape_dtype(t):
    """
    input: tensor
    output: 3 strings
    """
    t_numpy = t.cpu().numpy()
    t_bytes = str(t_numpy.tobytes())
    t_numpy_shape = str(t_numpy.shape)
    t_numpy_dtype = str(t_numpy.dtype)
    return t_bytes, t_numpy_shape, t_numpy_dtype

def analyze_image(img):
    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
    if cfg.DATASETS.TEST[0]=='icdar2019_test':
        md.set(thing_classes=["table"])
    else:
        md.set(thing_classes=["text","title","list","table","figure"])
    
    output = predictor(img)["instances"]
    v = Visualizer(img[:, :, ::-1],
                    md,
                    scale=1.0,
                    instance_mode=ColorMode.SEGMENTATION)
    result = v.draw_instance_predictions(output.to("cpu"))
    result_image = result.get_image()[:, :, ::-1]

    num_instances = len(output)
    image_size = output._image_size
    fields = list(output.get_fields().keys())
    for field in fields:
        if field == 'pred_boxes':
            boxes = output.get_fields()[field]
            boxes = boxes.tensor
            boxes_bytes, boxes_numpy_shape, boxes_numpy_dtype = get_bytes_shape_dtype(boxes)
            # boxes_recover = torch.from_numpy(np.frombuffer(boxes_bytes, dtype=boxes_numpy_dtype).reshape(boxes_numpy_shape))
        elif field == 'scores':
            scores = output.get_fields()[field]
            scores_bytes, scores_numpy_shape, scores_numpy_dtype = get_bytes_shape_dtype(scores)
            # scores_recover = torch.from_numpy(np.frombuffer(scores_bytes, dtype=scores_numpy_dtype).reshape(scores_numpy_shape))
        elif field == 'pred_classes':
            pred_classes = output.get_fields()[field]
            pred_classes_bytes, pred_classes_numpy_shape, pred_classes_numpy_dtype = get_bytes_shape_dtype(pred_classes)
            # pred_classes_recover = torch.from_numpy(np.frombuffer(pred_classes_bytes, dtype=pred_classes_numpy_dtype).reshape(pred_classes_numpy_shape))
    
    return result_image, num_instances, image_size, boxes_bytes, boxes_numpy_shape, boxes_numpy_dtype, scores_bytes, scores_numpy_shape, scores_numpy_dtype, pred_classes_bytes, pred_classes_numpy_shape, pred_classes_numpy_dtype
    
title = "Interactive demo: Document Layout Analysis with DiT"
description = "Demo for Microsoft's DiT, the Document Image Transformer for state-of-the-art document understanding tasks. This particular model is fine-tuned on PubLayNet, a large dataset for document layout analysis (read more at the links below). To use it, simply upload an image or use the example image below and click 'Submit'. Results will show up in a few seconds. If you want to make the output bigger, right-click on it and select 'Open image in new tab'."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2203.02378' target='_blank'>Paper</a> | <a href='https://github.com/microsoft/unilm/tree/master/dit' target='_blank'>Github Repo</a></p> | <a href='https://huggingface.co/docs/transformers/master/en/model_doc/dit' target='_blank'>HuggingFace doc</a></p>"
examples =[['publaynet_example.jpeg']]
css = ".output-image, .input-image, .image-preview {height: 600px !important}"

iface = gr.Interface(fn=analyze_image, 
                     inputs=gr.inputs.Image(type="numpy", label="document image"), 
                     outputs=[
                         gr.outputs.Image(type="numpy", label="annotated document"), 
                         gr.outputs.Textbox(label="num instances"),
                         gr.outputs.Textbox(label="image size (h,w in pixels)"),
                         gr.outputs.Textbox(label="boxes bytes"),
                         gr.outputs.Textbox(label="boxes numpy shape"),
                         gr.outputs.Textbox(label="boxes numpy dtype"),
                         gr.outputs.Textbox(label="scores bytes"),
                         gr.outputs.Textbox(label="scores numpy shape"),
                         gr.outputs.Textbox(label="scores numpy dtype"),
                         gr.outputs.Textbox(label="pred_classes bytes"),
                         gr.outputs.Textbox(label="pred_classes numpy shape"),
                         gr.outputs.Textbox(label="pred_classes numpy dtype")
                         ],
                     title=title,
                     description=description,
                     examples=examples,
                     article=article,
                     css=css
                     )
iface.launch(debug=True, cache_examples=True, enable_queue=True)