├── README.md └── external_masking.py /README.md: -------------------------------------------------------------------------------- 1 | ## Provides an external cv2 powered masking tool for https://github.com/AUTOMATIC1111/stable-diffusion-webui 2 | 3 | ## Installation 4 | **[Download the zipped script Here](https://github.com/dfaker/stable-diffusion-webui-cv2-external-masking-script/archive/refs/heads/main.zip)** 5 | and copy the file external_masking.py into your scripts folder. 6 | 7 | requires cv2 to be installed 8 | 9 | ```ShellSession 10 | py -m pip install opencv-python 11 | ``` 12 | 13 | Due to conflicting cv2 versions this extension cannot be used at the same time as the `d8ahazard/sd_dreambooth_extension` extension. 14 | 15 | ## Guide 16 | 17 | The UI inside stable-diffusion-webui is pretty simple 18 | ![Screenshot 2022-09-16 091930](https://user-images.githubusercontent.com/35278260/190592056-644c59db-907d-4cf1-ba85-0014eceea12a.jpg) 19 | 20 | `Masking preview size` controls the size of the popup CV2 window 21 | 22 | `Draw new mask on every run` will popup a new window for a new mask each time generate is clicked, usually it'll only appear on the first run, or when the input image is changed. 23 | 24 | The masking window itself is pretty minimal 25 | ![image](https://user-images.githubusercontent.com/35278260/193962552-3dfa4d28-5899-4e3f-a589-362de5990636.png) 26 | 27 | Showing the polygon currently being drawn in pink, left clicking starts a new polygon, right clicking closes the current polycon being drawn. 28 | 29 | C to the clear current mask. 30 | 31 | Q to quit and pass the current mask back to stable-diffusion-webui 32 | 33 | Scroll the mouse wheel to zoom in 34 | 35 | Middle click and drag to pan around the image 36 | 37 | The mask drawn with the script will not be shown on the input image, but will be used for all outputs: 38 | 39 | ![Screenshot 2022-09-16 091911](https://user-images.githubusercontent.com/35278260/190593109-10d47736-428c-4c3f-841a-a964778fbec7.jpg) 40 | 41 | ## Incompatible opencv versions 42 | 43 | Some users are reporting errors with the gui window functions like `highgui\src\window.cpp:1250: error: (-2:Unspecified error) The function is not implemented.` which seems to be down to having `opencv-python-headless` installed which doesn't include the gui code used to display the image windows, uninstall the current version and reinstall if you get a similar message: 44 | 45 | ```ShellSession 46 | # For global python: 47 | py -m pip uninstall opencv-python-headless 48 | py -m pip uninstall opencv-python 49 | py -m pip install --upgrade opencv-python 50 | # Or inside the stable-diffusion-webui venv: 51 | venv\Scripts\python -m pip uninstall opencv-python-headless 52 | venv\Scripts\python -m pip uninstall opencv-python 53 | venv\Scripts\python -m pip install --upgrade opencv-python 54 | ``` 55 | 56 | -------------------------------------------------------------------------------- /external_masking.py: -------------------------------------------------------------------------------- 1 | import math 2 | import os 3 | import sys 4 | import traceback 5 | 6 | 7 | import cv2 8 | from PIL import Image 9 | import numpy as np 10 | 11 | lastx,lasty=None,None 12 | zoomOrigin = 0,0 13 | zoomFactor = 1 14 | 15 | midDragStart = None 16 | 17 | def display_mask_ui(image,mask,max_size,initPolys): 18 | global lastx,lasty,zoomOrigin,zoomFactor 19 | 20 | lastx,lasty=None,None 21 | zoomOrigin = 0,0 22 | zoomFactor = 1 23 | 24 | polys = initPolys 25 | 26 | def on_mouse(event, x, y, buttons, param): 27 | global lastx,lasty,zoomFactor,midDragStart,zoomOrigin 28 | 29 | lastx,lasty = (x+zoomOrigin[0])/zoomFactor,(y+zoomOrigin[1])/zoomFactor 30 | 31 | if event == cv2.EVENT_LBUTTONDOWN: 32 | polys[-1].append((lastx,lasty)) 33 | elif event == cv2.EVENT_RBUTTONDOWN: 34 | polys.append([]) 35 | elif event == cv2.EVENT_MBUTTONDOWN: 36 | midDragStart = zoomOrigin[0]+x,zoomOrigin[1]+y 37 | elif event == cv2.EVENT_MBUTTONUP: 38 | if midDragStart is not None: 39 | zoomOrigin = max(0,midDragStart[0]-x),max(0,midDragStart[1]-y) 40 | midDragStart = None 41 | elif event == cv2.EVENT_MOUSEMOVE: 42 | if midDragStart is not None: 43 | zoomOrigin = max(0,midDragStart[0]-x),max(0,midDragStart[1]-y) 44 | elif event == cv2.EVENT_MOUSEWHEEL: 45 | origZoom = zoomFactor 46 | if buttons > 0: 47 | zoomFactor *= 1.1 48 | else: 49 | zoomFactor *= 0.9 50 | zoomFactor = max(1,zoomFactor) 51 | 52 | zoomOrigin = max(0,int(zoomOrigin[0]+ (max_size*0.25*(zoomFactor-origZoom)))) , max(0,int(zoomOrigin[1] + (max_size*0.25*(zoomFactor-origZoom)))) 53 | 54 | 55 | 56 | opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) 57 | 58 | if mask is None: 59 | opencvMask = cv2.cvtColor( np.array(opencvImage) , cv2.COLOR_BGR2GRAY) 60 | else: 61 | opencvMask = np.array(mask) 62 | 63 | 64 | maxdim = max(opencvImage.shape[1],opencvImage.shape[0]) 65 | 66 | factor = max_size/maxdim 67 | 68 | 69 | cv2.namedWindow('MaskingWindow', cv2.WINDOW_AUTOSIZE) 70 | cv2.setWindowProperty('MaskingWindow', cv2.WND_PROP_TOPMOST, 1) 71 | cv2.setMouseCallback('MaskingWindow', on_mouse) 72 | 73 | font = cv2.FONT_HERSHEY_SIMPLEX 74 | 75 | srcImage = opencvImage.copy() 76 | combinedImage = opencvImage.copy() 77 | 78 | interp = cv2.INTER_CUBIC 79 | if zoomFactor*factor < 0: 80 | interp = cv2.INTER_AREA 81 | 82 | zoomedSrc = cv2.resize(srcImage,(None,None),fx=zoomFactor*factor,fy=zoomFactor*factor,interpolation=interp) 83 | zoomedSrc = zoomedSrc[zoomOrigin[1]:zoomOrigin[1]+max_size,zoomOrigin[0]:zoomOrigin[0]+max_size,:] 84 | 85 | lastZoomFactor = zoomFactor 86 | lastZoomOrigin = zoomOrigin 87 | while 1: 88 | 89 | if lastZoomFactor != zoomFactor or lastZoomOrigin != zoomOrigin: 90 | interp = cv2.INTER_CUBIC 91 | if zoomFactor*factor < 0: 92 | interp = cv2.INTER_AREA 93 | zoomedSrc = cv2.resize(srcImage,(None,None),fx=zoomFactor*factor,fy=zoomFactor*factor,interpolation=interp) 94 | zoomedSrc = zoomedSrc[zoomOrigin[1]:zoomOrigin[1]+max_size,zoomOrigin[0]:zoomOrigin[0]+max_size,:] 95 | zoomedSrc = cv2.copyMakeBorder(zoomedSrc, 0, max_size-zoomedSrc.shape[0], 0, max_size-zoomedSrc.shape[1], cv2.BORDER_CONSTANT) 96 | 97 | lastZoomFactor = zoomFactor 98 | lastZoomOrigin = zoomOrigin 99 | 100 | foreground = np.zeros_like(zoomedSrc) 101 | 102 | for i,polyline in enumerate(polys): 103 | if len(polyline)>0: 104 | 105 | segs = polyline[::] 106 | 107 | active=False 108 | if len(polys[-1])>0 and i==len(polys)-1 and lastx is not None: 109 | segs = polyline+[(lastx,lasty)] 110 | active=True 111 | 112 | segs = np.array(segs) - np.array([(zoomOrigin[0]/zoomFactor,zoomOrigin[1]/zoomFactor)]) 113 | segs = (np.array([segs])*zoomFactor).astype(int) 114 | 115 | if active: 116 | cv2.fillPoly(foreground, (np.array(segs)) , ( 190, 107, 253), 0) 117 | else: 118 | cv2.fillPoly(foreground, (np.array(segs)) , (255, 255, 255), 0) 119 | 120 | if active: 121 | for x,y in segs[0]: 122 | cv2.circle(foreground, (int(x),int(y)), 5, (25,25,25), 3) 123 | cv2.circle(foreground, (int(x),int(y)), 5, (255,255,255), 2) 124 | 125 | 126 | foreground[foreground<1] = zoomedSrc[foreground<1] 127 | combinedImage = cv2.addWeighted(zoomedSrc, 0.5, foreground, 0.5, 0) 128 | 129 | helpText='Q=Save, C=Reset, LeftClick=Add new point to polygon, Rightclick=Close polygon, MouseWheel=Zoom, MidDrag=Pan' 130 | combinedImage = cv2.putText(combinedImage, helpText, (0,11), font, 0.4, (0,0,0), 2, cv2.LINE_AA) 131 | combinedImage = cv2.putText(combinedImage, helpText, (0,11), font, 0.4, (255,255,255), 1, cv2.LINE_AA) 132 | 133 | cv2.imshow('MaskingWindow',combinedImage) 134 | 135 | try: 136 | key = cv2.waitKey(1) 137 | if key == ord('q'): 138 | if len(polys[0])>0: 139 | newmask = np.zeros_like(cv2.cvtColor( opencvMask.astype('uint8') ,cv2.COLOR_GRAY2BGR) ) 140 | for i,polyline in enumerate(polys): 141 | if len(polyline)>0: 142 | segs = [(int(a/factor),int(b/factor)) for a,b in polyline] 143 | cv2.fillPoly(newmask, np.array([segs]), (255,255,255), 0) 144 | cv2.destroyWindow('MaskingWindow') 145 | return Image.fromarray( cv2.cvtColor( newmask, cv2.COLOR_BGR2GRAY) ),polys 146 | break 147 | if key == ord('c'): 148 | polys = [[]] 149 | 150 | except Exception as e: 151 | print(e) 152 | break 153 | 154 | cv2.destroyWindow('MaskingWindow') 155 | return mask,polys 156 | 157 | if __name__ == '__main__': 158 | img = Image.open('K:\\test2.png') 159 | oldmask = Image.new('L',img.size,(0,)) 160 | newmask,newPolys = display_mask_ui(img,oldmask,1024,[[]]) 161 | 162 | opencvImg = cv2.cvtColor( np.array(img) , cv2.COLOR_RGB2BGR) 163 | opencvMask = cv2.cvtColor( np.array(newmask) , cv2.COLOR_GRAY2BGR) 164 | 165 | combinedImage = cv2.addWeighted(opencvImg, 0.5, opencvMask, 0.5, 0) 166 | combinedImage = Image.fromarray( cv2.cvtColor( combinedImage , cv2.COLOR_BGR2RGB)) 167 | 168 | display_mask_ui(combinedImage,oldmask,1024,[[]]) 169 | 170 | 171 | exit() 172 | 173 | import modules.scripts as scripts 174 | import gradio as gr 175 | 176 | from modules.processing import Processed, process_images 177 | from modules.shared import opts, cmd_opts, state 178 | 179 | class Script(scripts.Script): 180 | 181 | def title(self): 182 | return "External Image Masking" 183 | 184 | def show(self, is_img2img): 185 | return is_img2img 186 | 187 | def ui(self, is_img2img): 188 | if not is_img2img: 189 | return None 190 | 191 | initialSize = 1024 192 | 193 | try: 194 | import tkinter as tk 195 | root = tk.Tk() 196 | screen_width = int(root.winfo_screenwidth()) 197 | screen_height = int(root.winfo_screenheight()) 198 | print(screen_width,screen_height) 199 | initialSize = min(screen_width,screen_height)-50 200 | print(initialSize) 201 | except Exception as e: 202 | print(e) 203 | 204 | max_size = gr.Slider(label="Masking preview size", minimum=512, maximum=initialSize*2, step=8, value=initialSize) 205 | with gr.Row(): 206 | ask_on_each_run = gr.Checkbox(label='Draw new mask on every run', value=False) 207 | non_contigious_split = gr.Checkbox(label='Process non-contigious masks separately', value=False) 208 | 209 | return [max_size,ask_on_each_run,non_contigious_split] 210 | 211 | def run(self, p, max_size, ask_on_each_run, non_contigious_split): 212 | 213 | if not hasattr(self,'lastImg'): 214 | self.lastImg = None 215 | 216 | if not hasattr(self,'lastMask'): 217 | self.lastMask = None 218 | 219 | if not hasattr(self,'lastPolys'): 220 | self.lastPolys = [[]] 221 | 222 | if ask_on_each_run or self.lastImg is None or self.lastImg != p.init_images[0]: 223 | 224 | if self.lastImg is None or self.lastImg != p.init_images[0]: 225 | self.lastPolys = [[]] 226 | 227 | p.image_mask,self.lastPolys = display_mask_ui(p.init_images[0],p.image_mask,max_size,self.lastPolys) 228 | self.lastImg = p.init_images[0] 229 | if p.image_mask is not None: 230 | self.lastMask = p.image_mask.copy() 231 | elif hasattr(self,'lastMask') and self.lastMask is not None: 232 | p.image_mask = self.lastMask.copy() 233 | 234 | if non_contigious_split: 235 | maskImgArr = np.array(p.image_mask) 236 | ret, markers = cv2.connectedComponents(maskImgArr) 237 | markerCount = markers.max() 238 | 239 | if markerCount > 1: 240 | tempimages = [] 241 | tempMasks = [] 242 | for maski in range(1,markerCount+1): 243 | print('maski',maski) 244 | maskSection = np.zeros_like(maskImgArr) 245 | maskSection[markers==maski] = 255 246 | p.image_mask = Image.fromarray( maskSection.copy() ) 247 | proc = process_images(p) 248 | images = proc.images 249 | tempimages.append(np.array(images[0])) 250 | tempMasks.append(np.array(maskSection.copy())) 251 | 252 | finalImage = tempimages[0].copy() 253 | 254 | for outimg,outmask in zip(tempimages,tempMasks): 255 | 256 | resizeimg = cv2.resize(outimg, (finalImage.shape[0],finalImage.shape[1]) ) 257 | resizedMask = cv2.resize(outmask, (finalImage.shape[0],finalImage.shape[1]) ) 258 | 259 | finalImage[resizedMask==255] = resizeimg[resizedMask==255] 260 | images = [finalImage] 261 | 262 | 263 | else: 264 | proc = process_images(p) 265 | images = proc.images 266 | else: 267 | proc = process_images(p) 268 | images = proc.images 269 | 270 | proc.images = images 271 | return proc 272 | --------------------------------------------------------------------------------