#!/usr/bin/env python3
#license: Apache-2.0
import subprocess
import sys
import os
target=(sys.argv[1:])

import os
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:,:,3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        # the float is how transparent
        color_mask = np.concatenate([np.random.random(3), [0.7]])
        img[m] = color_mask
    ax.imshow(img)

image = cv2.imread(target[0])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
size=image.shape[1]

sam_checkpoint = os.environ["SAM_PATH"]
model_type = os.environ["SAM_TYPE"]
device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

#mask_generator = SamAutomaticMaskGenerator(sam)
mask_generator = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,
    pred_iou_thresh=0.8,
    stability_score_thresh=0.9,
    crop_n_layers=1,
    crop_n_points_downscale_factor=2,
    min_mask_region_area=100,  # Requires open-cv to run post-processing
)
masks = mask_generator.generate(image)
len(masks)

plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')

masked_path=os.environ["SAM_OUTPUT"]
plt.savefig(masked_path+"masked.png", bbox_inches='tight', pad_inches = 0)

subprocess.run(["convert", masked_path+"masked.png", "-resize", str(size), masked_path+"masked_resized.png"])
subprocess.run(["mv", masked_path+"masked_resized.png", masked_path+"masked.png"])
# if show, it has to come last
# plt.show()