import modules.scripts as scripts
import gradio as gr
import os
import cv2
import numpy as np
from PIL import Image

from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
from scipy.ndimage import binary_erosion
from scripts.weshop_util import image_segment_by_mask


class Script(scripts.Script):  

# The title of the script. This is what will be displayed in the dropdown menu.
    def title(self):
        return "MaskSubject"


# Determines when the script should be shown in the dropdown menu via the 
# returned value. As an example:
# is_img2img is True if the current tab is img2img, and False if it is txt2img.
# Thus, return is_img2img to only show the script on the img2img tab.

    def show(self, is_img2img):
        return scripts.AlwaysVisible

# How the script's is displayed in the UI. See https://gradio.app/docs/#components
# for the different UI components you can use and how to create them.
# Most UI components can return a value, such as a boolean for a checkbox.
# The returned values are passed to the run method as parameters.

    def ui(self, is_img2img):
        useMaskSubject = gr.Checkbox(False, label="Use Mask Subject")
        return [useMaskSubject]

    def before_process(self, p, *args):
        pass

    def process(self, p, *args):
        """
        Author: jack
        This function is called before processing begins for AlwaysVisible scripts.
        You can modify the processing object (p) here, inject hooks, etc.
        args contains all values returned by components from ui()
        """
        useMaskSubject = args[0]

        # 扣主体
        if useMaskSubject and p.init_images is not None and len(p.init_images) > 0 and p.image_mask is not None:
            image = np.array(p.init_images[0])
            image_mask = np.array(p.image_mask)
            segment_result = image_segment_by_mask(image , image_mask)
            result = Image.fromarray(segment_result)
            p.init_images[0] = result
            print(f'process mask subject success\n')
  
    def postprocess(self, p, processed, *args):
        pass

# This is where the additional processing is implemented. The parameters include
# self, the model object "p" (a StableDiffusionProcessing class, see
# processing.py), and the parameters returned by the ui method.
# Custom functions can be defined here, and additional libraries can be imported 
# to be used in processing. The return value should be a Processed object, which is
# what is returned by the process_images method.

    # def run(self, p, angle, hflip, vflip, overwrite):
    #   pass