external_masking.py

import math
import os
import sys
import traceback


import cv2
from PIL import Image
import numpy as np

lastx,lasty=None,None
zoomOrigin = 0,0
zoomFactor = 1

midDragStart = None

def display_mask_ui(image,mask,max_size,initPolys):
  global lastx,lasty,zoomOrigin,zoomFactor

  lastx,lasty=None,None
  zoomOrigin = 0,0
  zoomFactor = 1

  polys = initPolys

  def on_mouse(event, x, y, buttons, param):
    global lastx,lasty,zoomFactor,midDragStart,zoomOrigin

    lastx,lasty = (x+zoomOrigin[0])/zoomFactor,(y+zoomOrigin[1])/zoomFactor

    if event == cv2.EVENT_LBUTTONDOWN:
      polys[-1].append((lastx,lasty))
    elif event == cv2.EVENT_RBUTTONDOWN:
      polys.append([])
    elif event == cv2.EVENT_MBUTTONDOWN:
        midDragStart = zoomOrigin[0]+x,zoomOrigin[1]+y
    elif event == cv2.EVENT_MBUTTONUP:
        if midDragStart is not None:
            zoomOrigin = max(0,midDragStart[0]-x),max(0,midDragStart[1]-y)
        midDragStart = None
    elif event == cv2.EVENT_MOUSEMOVE:
        if midDragStart is not None:
            zoomOrigin = max(0,midDragStart[0]-x),max(0,midDragStart[1]-y)
    elif event == cv2.EVENT_MOUSEWHEEL:
        origZoom = zoomFactor
        if buttons > 0:
            zoomFactor *= 1.1
        else:
            zoomFactor *= 0.9
        zoomFactor = max(1,zoomFactor)

        zoomOrigin = max(0,int(zoomOrigin[0]+ (max_size*0.25*(zoomFactor-origZoom)))) , max(0,int(zoomOrigin[1] + (max_size*0.25*(zoomFactor-origZoom))))



  opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

  if mask is None:
    opencvMask  = cv2.cvtColor( np.array(opencvImage) , cv2.COLOR_BGR2GRAY)
  else:
    opencvMask  = np.array(mask)


  maxdim = max(opencvImage.shape[1],opencvImage.shape[0])

  factor = max_size/maxdim


  cv2.namedWindow('MaskingWindow', cv2.WINDOW_AUTOSIZE)
  cv2.setWindowProperty('MaskingWindow', cv2.WND_PROP_TOPMOST, 1)
  cv2.setMouseCallback('MaskingWindow', on_mouse)

  font = cv2.FONT_HERSHEY_SIMPLEX

  srcImage = opencvImage.copy()
  combinedImage = opencvImage.copy()

  interp = cv2.INTER_CUBIC
  if zoomFactor*factor < 0:
    interp = cv2.INTER_AREA

  zoomedSrc = cv2.resize(srcImage,(None,None),fx=zoomFactor*factor,fy=zoomFactor*factor,interpolation=interp)
  zoomedSrc = zoomedSrc[zoomOrigin[1]:zoomOrigin[1]+max_size,zoomOrigin[0]:zoomOrigin[0]+max_size,:]

  lastZoomFactor = zoomFactor
  lastZoomOrigin = zoomOrigin
  while 1:

    if lastZoomFactor != zoomFactor or lastZoomOrigin != zoomOrigin:
        interp = cv2.INTER_CUBIC
        if zoomFactor*factor < 0:
          interp = cv2.INTER_AREA
        zoomedSrc = cv2.resize(srcImage,(None,None),fx=zoomFactor*factor,fy=zoomFactor*factor,interpolation=interp)
        zoomedSrc = zoomedSrc[zoomOrigin[1]:zoomOrigin[1]+max_size,zoomOrigin[0]:zoomOrigin[0]+max_size,:]
        zoomedSrc = cv2.copyMakeBorder(zoomedSrc, 0, max_size-zoomedSrc.shape[0], 0, max_size-zoomedSrc.shape[1], cv2.BORDER_CONSTANT)

        lastZoomFactor = zoomFactor
        lastZoomOrigin = zoomOrigin

    foreground    = np.zeros_like(zoomedSrc)

    for i,polyline in enumerate(polys):
      if len(polyline)>0:

        segs = polyline[::]

        active=False
        if len(polys[-1])>0 and i==len(polys)-1 and lastx is not None:
          segs = polyline+[(lastx,lasty)]
          active=True

        segs = np.array(segs) - np.array([(zoomOrigin[0]/zoomFactor,zoomOrigin[1]/zoomFactor)])
        segs = (np.array([segs])*zoomFactor).astype(int)
        
        if active:
          cv2.fillPoly(foreground, (np.array(segs)) , ( 190, 107,  253), 0)
        else:
          cv2.fillPoly(foreground, (np.array(segs)) , (255, 255, 255), 0)

        if active:
            for x,y in segs[0]:
                cv2.circle(foreground, (int(x),int(y)), 5, (25,25,25), 3)
                cv2.circle(foreground, (int(x),int(y)), 5, (255,255,255), 2)


    foreground[foreground<1] = zoomedSrc[foreground<1]
    combinedImage = cv2.addWeighted(zoomedSrc, 0.5, foreground, 0.5, 0)

    helpText='Q=Save, C=Reset, LeftClick=Add new point to polygon, Rightclick=Close polygon, MouseWheel=Zoom, MidDrag=Pan'
    combinedImage = cv2.putText(combinedImage, helpText, (0,11), font, 0.4, (0,0,0), 2, cv2.LINE_AA)
    combinedImage = cv2.putText(combinedImage, helpText, (0,11), font, 0.4, (255,255,255), 1, cv2.LINE_AA)

    cv2.imshow('MaskingWindow',combinedImage)

    try:
      key = cv2.waitKey(1)
      if key == ord('q'):
        if len(polys[0])>0:
          newmask = np.zeros_like(cv2.cvtColor( opencvMask.astype('uint8') ,cv2.COLOR_GRAY2BGR) )
          for i,polyline in enumerate(polys):
            if len(polyline)>0:
              segs = [(int(a/factor),int(b/factor)) for a,b in polyline]
              cv2.fillPoly(newmask, np.array([segs]), (255,255,255), 0)
          cv2.destroyWindow('MaskingWindow')
          return Image.fromarray( cv2.cvtColor( newmask, cv2.COLOR_BGR2GRAY) ),polys
        break
      if key == ord('c'):
        polys = [[]]

    except Exception as e:
      print(e)
      break

  cv2.destroyWindow('MaskingWindow')
  return mask,polys

if __name__ == '__main__':
    img  = Image.open('K:\\test2.png')
    oldmask = Image.new('L',img.size,(0,))
    newmask,newPolys = display_mask_ui(img,oldmask,1024,[[]])

    opencvImg  = cv2.cvtColor( np.array(img) , cv2.COLOR_RGB2BGR)
    opencvMask  = cv2.cvtColor( np.array(newmask) , cv2.COLOR_GRAY2BGR)

    combinedImage = cv2.addWeighted(opencvImg, 0.5, opencvMask, 0.5, 0)
    combinedImage = Image.fromarray( cv2.cvtColor(  combinedImage  , cv2.COLOR_BGR2RGB))
    
    display_mask_ui(combinedImage,oldmask,1024,[[]])


    exit()

import modules.scripts as scripts
import gradio as gr

from modules.processing import Processed, process_images
from modules.shared import opts, cmd_opts, state

class Script(scripts.Script):

    def title(self):
        return "External Image Masking"

    def show(self, is_img2img):
        return is_img2img

    def ui(self, is_img2img):
        if not is_img2img:
            return None

        initialSize = 1024

        try:
          import tkinter as tk
          root = tk.Tk()
          screen_width  = int(root.winfo_screenwidth())
          screen_height = int(root.winfo_screenheight())
          print(screen_width,screen_height)
          initialSize = min(screen_width,screen_height)-50
          print(initialSize)
        except Exception as e:
          print(e)

        max_size = gr.Slider(label="Masking preview size", minimum=512, maximum=initialSize*2, step=8, value=initialSize)
        with gr.Row():
          ask_on_each_run      = gr.Checkbox(label='Draw new mask on every run', value=False)
          non_contigious_split = gr.Checkbox(label='Process non-contigious masks separately', value=False)

        return [max_size,ask_on_each_run,non_contigious_split]

    def run(self, p, max_size, ask_on_each_run, non_contigious_split):

        if not hasattr(self,'lastImg'):
          self.lastImg = None

        if not hasattr(self,'lastMask'):
          self.lastMask = None

        if not hasattr(self,'lastPolys'):
          self.lastPolys = [[]]

        if ask_on_each_run or self.lastImg is None or self.lastImg != p.init_images[0]:

          if self.lastImg is None or self.lastImg != p.init_images[0]:
            self.lastPolys = [[]]

          p.image_mask,self.lastPolys  = display_mask_ui(p.init_images[0],p.image_mask,max_size,self.lastPolys)
          self.lastImg  = p.init_images[0]
          if p.image_mask is not None:
            self.lastMask = p.image_mask.copy()
        elif hasattr(self,'lastMask') and self.lastMask is not None:
          p.image_mask = self.lastMask.copy()

        if non_contigious_split:
          maskImgArr = np.array(p.image_mask)
          ret, markers = cv2.connectedComponents(maskImgArr)
          markerCount = markers.max()

          if markerCount > 1:
            tempimages = []
            tempMasks  = []
            for maski in range(1,markerCount+1):
              print('maski',maski)
              maskSection = np.zeros_like(maskImgArr)
              maskSection[markers==maski] = 255
              p.image_mask = Image.fromarray( maskSection.copy() )
              proc = process_images(p)
              images = proc.images
              tempimages.append(np.array(images[0]))
              tempMasks.append(np.array(maskSection.copy()))

            finalImage = tempimages[0].copy()

            for outimg,outmask in zip(tempimages,tempMasks):

              resizeimg = cv2.resize(outimg, (finalImage.shape[0],finalImage.shape[1]) )
              resizedMask = cv2.resize(outmask, (finalImage.shape[0],finalImage.shape[1]) )
              
              finalImage[resizedMask==255] = resizeimg[resizedMask==255]
            images = [finalImage]


          else:
            proc = process_images(p)
            images = proc.images
        else:
          proc = process_images(p)
          images = proc.images

        proc.images = images
        return proc