import os
import cv2
import sys
import numpy as np
import gradio as gr
from PIL import Image
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator


models = {
    'vit_b': './checkpoints/sam_vit_b_01ec64.pth',
    'vit_l': './checkpoints/sam_vit_l_0b3195.pth',
    'vit_h': './checkpoints/sam_vit_h_4b8939.pth'
}

def inference(device, model_type, input_img, points_per_side, pred_iou_thresh, stability_score_thresh, min_mask_region_area,
              stability_score_offset, box_nms_thresh, crop_n_layers, crop_nms_thresh):
    sam = sam_model_registry[model_type](checkpoint=models[model_type]).to(device)
    mask_generator = SamAutomaticMaskGenerator(
        sam,
        points_per_side=points_per_side,
        pred_iou_thresh=pred_iou_thresh,
        stability_score_thresh=stability_score_thresh,
        stability_score_offset=stability_score_offset,
        box_nms_thresh=box_nms_thresh,
        crop_n_layers=crop_n_layers,
        crop_nms_thresh=crop_nms_thresh,
        crop_overlap_ratio=512 / 1500,
        crop_n_points_downscale_factor=1,
        point_grids=None,
        min_mask_region_area=min_mask_region_area,
        output_mode='binary_mask'
    )

    masks = mask_generator.generate(input_img)
    sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True)

    mask_all = np.ones((input_img.shape[0], input_img.shape[1], 3))
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            mask_all[m==True, i] = color_mask[i]
    result = input_img / 255 * 0.3 + mask_all * 0.7

    return result, mask_all



with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown(
            '''# Segment Anything!🚀
            分割一切！CV的GPT-3时刻！
            [**官方网址**](https://segment-anything.com/)
            '''
        )
        with gr.Row():
            # 选择模型类型
            model_type = gr.Dropdown(["vit_b", "vit_l", "vit_h"], value='vit_b', label="选择模型")
            # 选择device
            device = gr.Dropdown(["cpu", "cuda"], value='cuda', label="选择你的硬件")

    # 参数
    with gr.Accordion(label='参数调整', open=False):
        with gr.Row():
            points_per_side = gr.Number(value=32, label="points_per_side", precision=0,
                                        info='''The number of points to be sampled along one side of the image. The total 
                                        number of points is points_per_side**2.''')
            pred_iou_thresh = gr.Slider(value=0.88, minimum=0, maximum=1.0, step=0.01, label="pred_iou_thresh",
                                        info='''A filtering threshold in [0,1], using the model's predicted mask quality.''')
            stability_score_thresh = gr.Slider(value=0.95, minimum=0, maximum=1.0, step=0.01, label="stability_score_thresh",
                                               info='''A filtering threshold in [0,1], using the stability of the mask under 
                                               changes to the cutoff used to binarize the model's mask predictions.''')
            min_mask_region_area = gr.Number(value=0, label="min_mask_region_area", precision=0,
                                             info='''If >0, postprocessing will be applied to remove disconnected regions 
                                             and holes in masks with area smaller than min_mask_region_area.''')
        with gr.Row():
            stability_score_offset = gr.Number(value=1, label="stability_score_offset",
                                               info='''The amount to shift the cutoff when calculated the stability score.''')
            box_nms_thresh = gr.Slider(value=0.7, minimum=0, maximum=1.0, step=0.01, label="box_nms_thresh",
                                       info='''The box IoU cutoff used by non-maximal ression to filter duplicate masks.''')
            crop_n_layers = gr.Number(value=0, label="crop_n_layers", precision=0,
                                      info='''If >0, mask prediction will be run again on crops of the image. 
                                      Sets the number of layers to run, where each layer has 2**i_layer number of image crops.''')
            crop_nms_thresh = gr.Slider(value=0.7, minimum=0, maximum=1.0, step=0.01, label="crop_nms_thresh",
                                        info='''The box IoU cutoff used by non-maximal suppression to filter duplicate 
                                        masks between different crops.''')

    # 显示图片
    with gr.Row().style(equal_height=True):
        with gr.Column():
            input_image = gr.Image(type="numpy")
            with gr.Row():
                button = gr.Button("Auto!")
        with gr.Tab(label='原图+mask'):
            image_output = gr.Image(type='numpy')
        with gr.Tab(label='Mask'):
            mask_output = gr.Image(type='numpy')

    gr.Examples(
        examples=[os.path.join(os.path.dirname(__file__), "./images/53960-scaled.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/2388455-scaled.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/1.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/2.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/3.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/4.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/5.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/6.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/7.jpg"),
                  os.path.join(os.path.dirname(__file__), "./images/8.jpg"),
                  ],
        inputs=input_image,
        outputs=image_output,
    )


    # 按钮交互
    button.click(inference, inputs=[device, model_type, input_image, points_per_side, pred_iou_thresh,
                                stability_score_thresh, min_mask_region_area, stability_score_offset, box_nms_thresh,
                                crop_n_layers, crop_nms_thresh],
             outputs=[image_output, mask_output])



demo.launch(debug=True)