import matplotlib.pyplot as plt
import matplotlib.pylab as pylab

import requests
from io import BytesIO
from PIL import Image
import numpy as np
pylab.rcParams['figure.figsize'] = 20, 12
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.engine.predictor_glip import GLIPDemo
import argparse
import pdb

def load(url_or_path):
    """
    Given an url or a path, this loads the file and
    """
    if url_or_path.startswith("http"):
        response = requests.get(url_or_path)
        pil_image = Image.open(BytesIO(response.content)).convert("RGB")
        # convert to BGR format
        image = np.array(pil_image)[:, :, [2, 1, 0]]
    else:
        image = np.array(Image.open(url_or_path).convert("RGB"))[:, :, [2, 1, 0]]
    return image


parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
parser.add_argument("--config", default="configs/pretrain/glip_Swin_T_O365_GoldG.yaml", metavar="FILE", help="path to config file", type=str)
parser.add_argument("--weight", default="OUTPUTS/GLIP_MODEL4/model_0020000.pth", metavar="FILE", help="path to weight file", type=str)
parser.add_argument("--image", default="http://farm4.staticflickr.com/3693/9472793441_b7822c00de_z.jpg", metavar="FILE", help="path to weight file", type=str)
parser.add_argument("--conf", default=0.4, type=float)
parser.add_argument("--caption", default="", type=str)
parser.add_argument("--ground_tokens", default=None, type=str)

args = parser.parse_args()

# update the config options with the config file
# manual override some options
cfg.local_rank = 0
cfg.num_gpus = 1
cfg.merge_from_file(args.config)
cfg.merge_from_list(["MODEL.WEIGHT", args.weight])
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])

glip_demo = GLIPDemo(
    cfg,
    min_image_size=800,
    show_mask_heatmaps=False
)

athetics_params = {
    "skip_name": False, # whether we overlay the phrase over the box
    "override_color": (255, 255, 255), # box color, default is white
    "text_size": 1.0,
    "text_pixel": 3,
    "box_alpha": 1.0,
    "box_pixel": 5,
    "text_offset_original": 8, # distance between text and box
}

image = load(args.image)
specified_tokens = args.ground_tokens.split(";") if args.ground_tokens is not None else None

result, _ = glip_demo.run_on_web_image(
    image,
    args.caption,
    args.conf,
    specified_tokens,
    **athetics_params)

plt.imshow(result[:, :, [2, 1, 0]])
plt.axis("off")
plt.savefig(args.image.replace('.png', "_demo.png").replace('.jpg', "_demo.jpg").replace('.jpeg', "_demo.jpeg"), bbox_inches='tight', pad_inches=0) # save as xxx_demo.xxx