├── .github └── ISSUE_TEMPLATE │ ├── feature_request.md │ └── bug_report.md ├── style.css ├── README.md ├── javascript └── training_picker.js └── scripts └── training_picker.py /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: enhancement 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | Please cover these three things in your report: 11 | * What exactly is wrong 12 | * How to recreate the problem 13 | * What you expected to happen vs what actually happened 14 | 15 | You can also optionally provide these details if you think it'll help me figure out the issue faster: 16 | * Operating system 17 | * Browser 18 | * Any error messages or logs, if applicable 19 | * Particular details about how you installed the extension or other relevant software 20 | * Screenshots demonstrating the issue 21 | -------------------------------------------------------------------------------- /style.css: -------------------------------------------------------------------------------- 1 | #cropPreviewRect { 2 | position:absolute; 3 | top:0px; 4 | left:0px; 5 | border:2px solid red; 6 | background:rgba(255, 0, 0, 0.3); 7 | z-index: 900; 8 | pointer-events:none; 9 | display:none 10 | } 11 | 12 | #crop_options_row { 13 | max-height: 60px; 14 | } 15 | 16 | #crop_preview { 17 | max-height: 400px; 18 | } 19 | 20 | #crop_parameters { 21 | display: none; 22 | } 23 | 24 | #frame_number * { 25 | border: none; 26 | background: none; 27 | padding: 0px; 28 | } 29 | 30 | #frame_number input { 31 | float: right 32 | } 33 | 34 | [id^=refresh_], [id^=open_folder_] { 35 | display: flex; 36 | align-content: center; 37 | justify-content: center; 38 | max-width: 2.5em; 39 | min-width: 2.5em; 40 | height: 2.4em; 41 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # training-picker 2 | 3 | Adds a tab to the webui that allows the user to automatically extract keyframes from video, and manually extract 512x512 crops of those frames for use in model training. 4 | 5 | ![image](https://user-images.githubusercontent.com/2313721/200236386-5fed34df-03e4-4ea6-a653-e1b60393afcd.png) 6 | 7 | ## Installation: 8 | 9 | 1. Install [AUTOMATIC1111's Stable Diffusion Webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) 10 | 2. Install [ffmpeg](https://ffmpeg.org/) for your operating system 11 | 3. Clone this repository into the `extensions` folder inside the webui 12 | 13 | Make sure you don't already have `python-ffmpeg` installed globally, as the library this program uses is `ffmpeg-python`, and your installation will conflict with it. 14 | 15 | ## Usage: 16 | 17 | ### Creating an extracted frame set 18 | 19 | 1. Drop videos you want to extract cropped frames from into the `training-picker/videos` folder 20 | 2. Open up the Training Picker tab of the webui 21 | 3. Select one of the videos you placed in the `training-picker/videos` folder from the dropdown on the left 22 | 4. Click 'Extract Frames` 23 | 5. After the extraction finishes, a new keyframe set for the video should be selectable in the dropdown on the right 24 | 6. Select the keyframe set, and the frames will appear in the browser below 25 | 26 | Optionally, you can also just supply a large collection of individual images you would like to work with directly by placing them into a folder within `training-picker/extracted-frames`. 27 | 28 | ### Cropping 29 | 30 | * Scroll up and down to increase or decrease the size of the crop brush 31 | * Ctrl+scroll to adjust the aspect ratio of the crop, and middle-click to reset the aspect ratio to 1:1 32 | * Shift+scroll to adjust the size / aspect ratio by smaller increments 33 | * Click to save a crop at the brush's position 34 | * Navigate between frames in the collection by clicking the navigation buttons, entering a number into the counter, or by using the arrow keys / AD 35 | * Select an outfill method to outfill the non-square area of a rectangular crop into a square shape 36 | * Click "Bulk process frames with chosen outfill method" to automatically process every image in the current frame set using the outfill method chosen, outputting to the directory under "Save crops to:" 37 | * Crops will be saved to `training-picker/cropped-frames` by default 38 | -------------------------------------------------------------------------------- /javascript/training_picker.js: -------------------------------------------------------------------------------- 1 | const intervalStep = 8; 2 | 3 | let centerSize = 512; 4 | let aspectRatio = 0; 5 | let mouseOn = false; 6 | let mXPos = 0; 7 | let mYPos = 0; 8 | let xPos = 0; 9 | let yPos = 0; 10 | 11 | function cropPreviewRect() { 12 | var cropPreviewRect = gradioApp().querySelector("#cropPreviewRect"); 13 | if (!cropPreviewRect) { 14 | cropPreviewRect = document.createElement("div"); 15 | cropPreviewRect.id = "cropPreviewRect"; 16 | gradioApp().appendChild(cropPreviewRect); 17 | } 18 | return cropPreviewRect; 19 | } 20 | 21 | function getBrushDim() { 22 | let img = gradioApp().querySelector("#frame_browser img"); 23 | let bound = img.getBoundingClientRect(); 24 | let ar = Math.exp(aspectRatio); 25 | let brushW = centerSize / ar; 26 | let brushH = centerSize * ar; 27 | brushW = Math.min(Math.max(intervalStep, brushW), bound.width); 28 | brushH = Math.min(Math.max(intervalStep, brushH), bound.height); 29 | return [brushW, brushH]; 30 | } 31 | 32 | function updateCpr() { 33 | let img = gradioApp().querySelector("#frame_browser img"); 34 | let bound = img.getBoundingClientRect(); 35 | let [brushW, brushH] = getBrushDim(); 36 | xPos = Math.min(Math.max(bound.left + brushW / 2 + window.scrollX, mXPos), bound.right - brushW / 2 + window.scrollX); 37 | yPos = Math.min(Math.max(bound.top + brushH / 2 + window.scrollY, mYPos), bound.bottom - brushH / 2 + window.scrollY); 38 | let cpr = cropPreviewRect(); 39 | cpr.style.width = brushW + "px"; 40 | cpr.style.height = brushH + "px"; 41 | cpr.style.left = (xPos - brushW / 2) + "px"; 42 | cpr.style.top = (yPos - brushH / 2) + "px"; 43 | cpr.style.display = mouseOn ? 'block' : 'none'; 44 | } 45 | 46 | function resetAspectRatio() { 47 | aspectRatio = 0; 48 | updateCpr(); 49 | } 50 | 51 | document.addEventListener("DOMContentLoaded", () => { 52 | (new MutationObserver(e => { 53 | let img = gradioApp().querySelector("#frame_browser img"); 54 | if (img && (img.getAttribute("listeners") !== "true")) { 55 | img.addEventListener("mousemove", e => { 56 | mXPos = e.pageX; 57 | mYPos = e.pageY; 58 | updateCpr(); 59 | }); 60 | img.addEventListener("mouseenter", e => { 61 | mouseOn = true; 62 | updateCpr(); 63 | }); 64 | img.addEventListener("mouseleave", e => { 65 | mouseOn = false; 66 | updateCpr(); 67 | }); 68 | img.addEventListener("wheel", e => { 69 | e.preventDefault(); 70 | let img = gradioApp().querySelector("#frame_browser img"); 71 | let bound = img.getBoundingClientRect(); 72 | let x = e.deltaY/100; 73 | if (e.shiftKey) x /= 8; 74 | if (e.ctrlKey) { 75 | aspectRatio -= x * 0.05; 76 | aspectRatio = Math.max(-2, Math.min(aspectRatio, 2)) 77 | } else { 78 | centerSize -= x * intervalStep; 79 | centerSize = Math.max(intervalStep, Math.min(centerSize, Math.max(bound.width, bound.height))) 80 | } 81 | updateCpr(); 82 | }); 83 | img.addEventListener("mousedown", e => { 84 | if (e.button == 0) { 85 | updateCpr(); 86 | let xRatio = img.naturalWidth / img.width; 87 | let yRatio = img.naturalHeight / img.height; 88 | let bound = img.getBoundingClientRect(); 89 | let [brushW, brushH] = getBrushDim(); 90 | let cropData = { 91 | x1: Math.floor((xPos - brushW / 2 - bound.left - window.scrollX) * xRatio), 92 | y1: Math.floor((yPos - brushH / 2 - bound.top - window.scrollY) * yRatio), 93 | x2: Math.floor((xPos + brushW / 2 - bound.left - window.scrollX) * xRatio), 94 | y2: Math.floor((yPos + brushH / 2 - bound.top - window.scrollY) * yRatio) 95 | }; 96 | let crop_parameters = gradioApp().querySelector("#crop_parameters textarea"); 97 | crop_parameters.value = JSON.stringify(cropData); 98 | crop_parameters.dispatchEvent(new CustomEvent("input", {})); // necessary to notify gradio that the value has changed 99 | gradioApp().querySelector("#crop_button").click(); 100 | } else if (e.button == 1) { 101 | resetAspectRatio(); 102 | } 103 | }); 104 | img.setAttribute("listeners", "true"); 105 | } 106 | })).observe(gradioApp(), {childList: true, subtree: true}); 107 | 108 | document.addEventListener("keydown", e => { 109 | if (mouseOn) { 110 | if (e.code == "ArrowRight" || e.code == "KeyD") { 111 | gradioApp().querySelector("#next_button").click(); 112 | } else if (e.code == "ArrowLeft" || e.code == "KeyA") { 113 | gradioApp().querySelector("#prev_button").click(); 114 | } 115 | } 116 | }); 117 | }); 118 | -------------------------------------------------------------------------------- /scripts/training_picker.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import platform 3 | import math 4 | import json 5 | import sys 6 | import os 7 | import re 8 | from pathlib import Path 9 | 10 | import gradio as gr 11 | import numpy as np 12 | from tqdm import tqdm 13 | from PIL import Image, ImageFilter 14 | import cv2 15 | 16 | from modules.ui import create_refresh_button 17 | from modules.ui_common import folder_symbol 18 | from modules.shared import opts, OptionInfo 19 | from modules import shared, paths, script_callbacks 20 | 21 | current_frame_set = [] 22 | current_frame_set_index = 0 23 | 24 | class CachedImage: 25 | def __init__(self, path): 26 | self.path = path 27 | self.image = None 28 | 29 | def get(self): 30 | if self.image == None: 31 | self.image = Image.open(self.path) 32 | return self.image 33 | 34 | # copied directly from ui.py 35 | # if it were not defined inside another function, i would import it instead 36 | def open_folder(f): 37 | if not os.path.exists(f): 38 | print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.') 39 | return 40 | elif not os.path.isdir(f): 41 | print(f""" 42 | WARNING 43 | An open_folder request was made with an argument that is not a folder. 44 | This could be an error or a malicious attempt to run code on your computer. 45 | Requested path was: {f} 46 | """, file=sys.stderr) 47 | return 48 | 49 | if not shared.cmd_opts.hide_ui_dir_config: 50 | path = os.path.normpath(f) 51 | if platform.system() == "Windows": 52 | os.startfile(path) 53 | elif platform.system() == "Darwin": 54 | subprocess.Popen(["open", path]) 55 | else: 56 | subprocess.Popen(["xdg-open", path]) 57 | 58 | def create_open_folder_button(path, elem_id): 59 | button = gr.Button(folder_symbol, elem_id=elem_id) 60 | if 'gradio.templates' in getattr(path, "__module__", ""): 61 | button.click(fn=lambda p: open_folder(p), inputs=[path], outputs=[]) 62 | else: 63 | button.click(fn=lambda: open_folder(path), inputs=[], outputs=[]) 64 | return button 65 | 66 | def resized_background(im, color): 67 | dim = max(*im.size) 68 | w, h = im.size 69 | background = Image.new(mode="RGBA", size=(dim, dim), color=color) 70 | background.paste(im, (dim // 2 - w // 2, dim // 2 - h // 2)) 71 | return background 72 | 73 | def gradient_blur(im, factor, original_dims): 74 | w, h = original_dims 75 | nw, nh = im.size 76 | n = abs(w - h) // 2 77 | original = im.copy() 78 | if w > h: 79 | for y in range(n): 80 | top_sliver = (0, n - y - 1, w, n - y) 81 | bottom_sliver = (0, (nh + h) // 2 + y + 1, w, (nh + h) // 2 + y + 2) 82 | blurred = original.filter(ImageFilter.GaussianBlur(factor * (y/n))) 83 | im.paste(blurred.crop(top_sliver), top_sliver) 84 | im.paste(blurred.crop(bottom_sliver), bottom_sliver) 85 | else: 86 | for x in range(n): 87 | left_sliver = (n - x - 1, 0, n - x, h) 88 | right_sliver = ((nw + w) // 2 + x, 0, (nw + w) // 2 + x + 1, h) 89 | blurred = original.filter(ImageFilter.GaussianBlur(factor * (x/n))) 90 | im.paste(blurred.crop(left_sliver), left_sliver) 91 | im.paste(blurred.crop(right_sliver), right_sliver) 92 | return im 93 | 94 | # Outfill methods 95 | 96 | def no_outfill(im, **kwargs): 97 | return im 98 | 99 | def stretch(im, **kwargs): 100 | dim = max(*im.size) 101 | return im.resize((dim, dim)) 102 | 103 | def transparent(im, **kwargs): 104 | return resized_background(im, (0,0,0,0)) 105 | 106 | def solid(im, **kwargs): 107 | return resized_background(im, kwargs['color']) 108 | 109 | def average(im, **kwargs): 110 | return resized_background(im, tuple(int(x) for x in np.asarray(im).mean(axis=0).mean(axis=0))) 111 | 112 | def dominant(im, **kwargs): 113 | _, labels, palette = cv2.kmeans( 114 | np.float32(np.asarray(im).reshape(-1, 3)), 115 | kwargs['n_clusters'], 116 | None, 117 | (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.1), 118 | 10, 119 | cv2.KMEANS_RANDOM_CENTERS 120 | ) 121 | _, counts = np.unique(labels, return_counts=True) 122 | color = palette[np.argmax(counts)] 123 | return resized_background(im, tuple(int(x) for x in color)) 124 | 125 | def border_stretch(im, **kwargs): 126 | w, h = im.size 127 | fw, fh = kwargs['dim_override'] if 'dim_override' in kwargs else [max(w, h)]*2 128 | arr = np.asarray(im) 129 | axis = -1 130 | if 'axis_override' in kwargs: 131 | axis = kwargs['axis_override'] 132 | elif w > h: 133 | axis = 0 134 | elif h > w: 135 | axis = 1 136 | if axis == 0: 137 | n = abs(fh - h) // 2 138 | o = 1 139 | arr = np.repeat(arr, [n + o] + [1]*(h-1), axis=0) 140 | if (fh - h) % 2 == 1: 141 | o += 1 142 | arr = np.repeat(arr, [1]*(h+n-1) + [n + o], axis=0) 143 | elif axis == 1: 144 | n = abs(fw - w) // 2 145 | o = 1 146 | arr = np.repeat(arr, [n + o] + [1]*(w-1), axis=1) 147 | if (fw - w) % 2 == 1: 148 | o += 1 149 | arr = np.repeat(arr, [1]*(w+n-1) + [n + o], axis=1) 150 | final = Image.fromarray(arr) 151 | if kwargs['blur'] > 0: 152 | final = gradient_blur(final, kwargs['blur'], (w, h)) 153 | return final 154 | 155 | def reflect(im, **kwargs): 156 | w, h = im.size 157 | fw, fh = kwargs['dim_override'] if 'dim_override' in kwargs else [max(w, h)]*2 158 | arr = np.asarray(im) 159 | base = arr.copy() 160 | axis = -1 161 | if 'axis_override' in kwargs: 162 | axis = kwargs['axis_override'] 163 | elif w > h: 164 | axis = 0 165 | elif h > w: 166 | axis = 1 167 | if axis == 0: 168 | while arr.shape[0] < fh: 169 | arr = np.concatenate((arr, base[::-1,:])) 170 | arr = np.concatenate((base[::-1,:], arr)) 171 | base = base[::-1,:] 172 | n = abs(arr.shape[0] - fh) // 2 173 | arr = arr[n:-n, :] 174 | elif axis == 1: 175 | while arr.shape[1] < fw: 176 | arr = np.concatenate((arr, base[:,::-1]), axis=1) 177 | arr = np.concatenate((base[:,::-1], arr), axis=1) 178 | base = base[:,::-1] 179 | n = abs(fw - arr.shape[1]) // 2 180 | arr = arr[:, n:-n] 181 | final = Image.fromarray(arr) 182 | if kwargs['blur'] > 0: 183 | final = gradient_blur(final, kwargs['blur'], (w, h)) 184 | return final 185 | 186 | def blur(im, **kwargs): 187 | dim = max(*im.size) 188 | w, h = im.size 189 | background = im.resize((dim, dim)) 190 | background = background.filter(ImageFilter.GaussianBlur(kwargs['blur'])) 191 | background.paste(im, (dim // 2 - w // 2, dim // 2 - h // 2)) 192 | return background 193 | 194 | def keep_original(im, **kwargs): 195 | final = kwargs['original'] 196 | if kwargs['blur'] > 0: 197 | final = gradient_blur(final, kwargs['blur'], im.size) 198 | return final 199 | 200 | outfill_methods = { 201 | "Don't outfill": no_outfill, 202 | "Stretch image": stretch, 203 | "Transparent": transparent, 204 | "Solid color": solid, 205 | "Average image color": average, 206 | "Dominant image color": dominant, 207 | "Stretch pixels at border": border_stretch, 208 | "Reflect image around border": reflect, 209 | "Blurred & stretched overlay": blur, 210 | "Reuse original image": keep_original 211 | } 212 | 213 | def on_ui_tabs(): 214 | 215 | fixed_size = int(opts.training_picker_fixed_size) 216 | videos_path = Path(opts.training_picker_videos_path) 217 | framesets_path = Path(opts.training_picker_framesets_path) 218 | default_output_path = Path(opts.training_picker_default_output_path) 219 | 220 | for p in [videos_path, framesets_path]: 221 | os.makedirs(p, exist_ok=True) 222 | 223 | def get_videos_list(): 224 | return list(v.name for v in videos_path.iterdir() if v.suffix in [".mp4", ".mkv", ".webm", ".gif"]) 225 | 226 | def get_framesets_list(): 227 | return list(v.name for v in framesets_path.iterdir() if v.is_dir()) 228 | 229 | with gr.Blocks(analytics_enabled=False) as training_picker: 230 | videos_list = get_videos_list() 231 | framesets_list = get_framesets_list() 232 | 233 | # structure 234 | with gr.Row(): 235 | with gr.Column(): 236 | with gr.Row(): 237 | video_dropdown = gr.Dropdown(choices=videos_list, elem_id="video_dropdown", label="Video to extract frames from:") 238 | create_refresh_button(video_dropdown, lambda: None, lambda: {"choices": get_videos_list()}, "refresh_videos_list") 239 | create_open_folder_button(videos_path, "open_folder_videos") 240 | with gr.Row(): 241 | only_keyframes_checkbox = gr.Checkbox(value=True, label="Only extract keyframes (recommended)") 242 | with gr.Column(visible=False) as frame_skip_container: 243 | frame_skip_input = gr.Number(value=1, min=1, label="Extract every nth frame", interactive=True) 244 | extract_frames_button = gr.Button(value="Extract Frames", variant="primary") 245 | log_output = gr.HTML(value="") 246 | with gr.Column(): 247 | with gr.Row(): 248 | frameset_dropdown = gr.Dropdown(choices=framesets_list, elem_id="frameset_dropdown", label="Extracted Frame Set", interactive=True) 249 | create_refresh_button(frameset_dropdown, lambda: None, lambda: {"choices": get_framesets_list()}, "refresh_framesets_list") 250 | create_open_folder_button(framesets_path, "open_folder_framesets") 251 | with gr.Row(elem_id="crop_options_row"): 252 | resize_checkbox = gr.Checkbox(value=True, label=f"Resize crops to {fixed_size}x{fixed_size}") 253 | with gr.Row(): 254 | outfill_setting = gr.Dropdown(choices=list(outfill_methods.keys()), value="Don't outfill", label="Outfill method:", interactive=True) 255 | with gr.Row(): 256 | with gr.Column(): 257 | reset_aspect_ratio_button = gr.Button(value="Reset Aspect Ratio") 258 | bulk_process_button = gr.Button(value="Process", variant="primary") 259 | with gr.Row(visible=False) as outfill_setting_options: 260 | with gr.Column(visible=False, scale=0.3) as original_image_outfill_setting_container: 261 | outfill_original_image_outfill_setting = gr.Dropdown(label="Image border outfill method:", scale=0.3, value="Stretch pixels at border", choices=["Stretch pixels at border", "Reflect image around border", "Black outfill"]) 262 | with gr.Column(visible=False) as color_container: 263 | outfill_color = gr.ColorPicker(value="#000000", label="Outfill border color:", interactive=True) 264 | with gr.Column(visible=False) as border_blur_container: 265 | outfill_border_blur = gr.Slider(value=0, min=0, max=100, step=0.01, label="Blur amount:", interactive=True) 266 | with gr.Column(visible=False) as n_clusters_container: 267 | outfill_n_clusters = gr.Slider(value=5, min=1, max=50, step=1, label="Number of clusters:", interactive=True) 268 | with gr.Row(): 269 | output_dir = gr.Text(value=default_output_path, label="Save crops to:") 270 | create_open_folder_button(output_dir, "open_folder_crops") 271 | with gr.Row(): 272 | with gr.Column(): 273 | crop_preview = gr.Image(interactive=False, elem_id="crop_preview", show_label=False) 274 | with gr.Column(): 275 | frame_browser = gr.Image(interactive=False, elem_id="frame_browser", show_label=False) 276 | with gr.Row(): 277 | prev_button = gr.Button(value="<", elem_id="prev_button") 278 | with gr.Row(): 279 | frame_number = gr.Number(value=0, elem_id="frame_number", live=True, show_label=False) 280 | frame_max = gr.HTML(value="", elem_id="frame_max") 281 | next_button = gr.Button(value=">", elem_id="next_button") 282 | 283 | # invisible elements 284 | crop_button = gr.Button(elem_id="crop_button", visible=False) 285 | crop_parameters = gr.Text(elem_id="crop_parameters", visible=False) 286 | 287 | # events 288 | def extract_frames_button_click(video_file, only_keyframes, frame_skip_input): 289 | try: 290 | import ffmpeg 291 | except ModuleNotFoundError: 292 | print("Installing ffmpeg-python") 293 | subprocess.check_call([sys.executable, "-m", "pip", "install", "ffmpeg-python"]) 294 | import ffmpeg 295 | try: 296 | print(f"Extracting frames from {video_file}") 297 | full_path = videos_path / video_file 298 | output_path = framesets_path / Path(video_file).stem 299 | if output_path.is_dir(): 300 | print("Directory already exists!") 301 | return gr.update(), f"Frame set already exists at {output_path}! Delete the folder first if you would like to recreate it." 302 | os.makedirs(output_path, exist_ok=True) 303 | output_name_fmat = str((output_path / "%02d.png").resolve()) 304 | if only_keyframes: 305 | stream = ffmpeg.input( 306 | str(full_path), 307 | skip_frame="nokey", 308 | vsync="vfr", 309 | ) 310 | stream.output(output_name_fmat).run() 311 | else: 312 | stream = ffmpeg.input( 313 | str(full_path), 314 | vsync="vfr", 315 | ) 316 | stream.output( 317 | output_name_fmat, 318 | vf=f"select=not(mod(n\,{frame_skip_input}))", 319 | ).run() 320 | print("Extraction complete!") 321 | return gr.Dropdown.update(choices=get_framesets_list()), f"Successfully created frame set {output_path.name}" 322 | except Exception as e: 323 | print(f"Exception encountered while attempting to extract frames: {e}") 324 | return gr.update(), f"Error: {e}" 325 | extract_frames_button.click(fn=extract_frames_button_click, inputs=[video_dropdown, only_keyframes_checkbox, frame_skip_input], outputs=[frameset_dropdown, log_output]) 326 | 327 | def get_image_update(): 328 | global current_frame_set_index 329 | global current_frame_set 330 | return gr.Image.update(value=current_frame_set[current_frame_set_index].get()), current_frame_set_index+1, f"/{len(current_frame_set)}" 331 | 332 | def null_image_update(): 333 | return gr.update(), 0, "" 334 | 335 | def only_keyframes_checkbox_change(only_keyframes_checkbox): 336 | return gr.update(visible=not only_keyframes_checkbox) 337 | only_keyframes_checkbox.change(fn=only_keyframes_checkbox_change, inputs=[only_keyframes_checkbox], outputs=[frame_skip_container]) 338 | 339 | def frameset_dropdown_change(frameset): 340 | global current_frame_set_index 341 | global current_frame_set 342 | current_frame_set_index = 0 343 | full_path = framesets_path / frameset 344 | current_frame_set = [CachedImage(impath) for impath in full_path.iterdir() if impath.suffix in [".png", ".jpg"]] 345 | try: current_frame_set = sorted(current_frame_set, key=lambda f:int(re.match(r"^(\d+).*", f.path.name).group(1))) 346 | except Exception as e: print(f"Unable to sort frames: {e}") 347 | return get_image_update() 348 | frameset_dropdown.change(fn=frameset_dropdown_change, inputs=[frameset_dropdown], outputs=[frame_browser, frame_number, frame_max]) 349 | 350 | def prev_button_click(): 351 | global current_frame_set_index 352 | global current_frame_set 353 | if current_frame_set != []: 354 | current_frame_set_index = (current_frame_set_index - 1) % len(current_frame_set) 355 | return get_image_update() 356 | return null_image_update() 357 | prev_button.click(fn=prev_button_click, inputs=[], outputs=[frame_browser, frame_number, frame_max]) 358 | 359 | def next_button_click(): 360 | global current_frame_set_index 361 | global current_frame_set 362 | if current_frame_set != []: 363 | current_frame_set_index = (current_frame_set_index + 1) % len(current_frame_set) 364 | return get_image_update() 365 | return null_image_update() 366 | next_button.click(fn=next_button_click, inputs=[], outputs=[frame_browser, frame_number, frame_max]) 367 | 368 | def frame_number_change(frame_number): 369 | global current_frame_set_index 370 | global current_frame_set 371 | if current_frame_set != []: 372 | current_frame_set_index = int(min(max(0, frame_number - 1), len(current_frame_set) - 1)) 373 | return get_image_update() 374 | return null_image_update() 375 | frame_number.change(fn=frame_number_change, inputs=[frame_number], outputs=[frame_browser, frame_number, frame_max]) 376 | 377 | def process_image(image, should_resize, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters, square_original): 378 | w, h = image.size 379 | if should_resize: 380 | ratio = fixed_size / max(w, h) 381 | image = image.resize((math.ceil(w * ratio), math.ceil(h * ratio))) 382 | if square_original: 383 | square_original = square_original.resize((fixed_size - 1, fixed_size - 1)) # i would prefer to resize to the exact fixed size but a sliver of unblurred image appears otherwise in the final result :/ 384 | if outfill_setting != "Don't outfill": 385 | image = outfill_methods[outfill_setting](image, color=outfill_color, blur=outfill_border_blur, n_clusters=outfill_n_clusters, original=square_original) 386 | return image 387 | 388 | def get_squared_original(full_im, bounds, outfill_method): 389 | x1, y1, x2, y2 = bounds 390 | w, h = x2 - x1, y2 - y1 391 | cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 392 | r = max(w, h) // 2 393 | iw, ih = full_im.size 394 | outrad = max(iw, ih) 395 | dim_override = (int(outrad*2), int(outrad*2)) 396 | ox, oy = (0, 0) if outfill_method == "Black outfill" else (outrad // 2 + (outrad - iw) // 2, outrad // 2 + (outrad - ih) // 2) 397 | new_bounds = (cx - r + ox, cy - r + oy, cx + r + ox, cy + r + oy) 398 | if outfill_method == "Stretch pixels at border": 399 | full_im = border_stretch(full_im, blur=0, dim_override=dim_override, axis_override=0) 400 | full_im = border_stretch(full_im, blur=0, dim_override=dim_override, axis_override=1) 401 | elif outfill_method == "Reflect image around border": 402 | full_im = reflect(full_im, blur=0, dim_override=dim_override, axis_override=0) 403 | full_im = reflect(full_im, blur=0, dim_override=dim_override, axis_override=1) 404 | return full_im.crop(new_bounds) 405 | 406 | def crop_button_click(raw_params, frame_browser, should_resize, output_dir, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters, outfill_original_image_outfill_setting): 407 | params = json.loads(raw_params) 408 | im = Image.fromarray(frame_browser) 409 | crop_boundary = (params['x1'], params['y1'], params['x2'], params['y2']) 410 | cropped = im.crop(crop_boundary) 411 | if outfill_setting == "Reuse original image": 412 | square_original = get_squared_original(im, crop_boundary, outfill_original_image_outfill_setting) 413 | else: 414 | square_original = None 415 | cropped = process_image(cropped, should_resize, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters, square_original) 416 | save_path = Path(output_dir) 417 | os.makedirs(str(save_path.resolve()), exist_ok=True) 418 | current_images = [r for r in (re.match(r"(\d+).png", f.name) for f in save_path.iterdir()) if r] 419 | if current_images == []: 420 | next_image_num = 0 421 | else: 422 | next_image_num = 1 + max(int(r.group(1)) for r in current_images) 423 | filename = save_path / f"{next_image_num}.png" 424 | cropped.save(filename) 425 | return gr.Image.update(value=cropped), f"Saved to {filename}" 426 | crop_button.click(fn=crop_button_click, inputs=[crop_parameters, frame_browser, resize_checkbox, output_dir, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters, outfill_original_image_outfill_setting], outputs=[crop_preview, log_output]) 427 | 428 | def bulk_process_button_click(frameset, should_resize, output_dir, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters): 429 | if outfill_setting == "Reuse original image": 430 | return gr.Image.update(value="https://user-images.githubusercontent.com/2313721/200725535-d2aca52a-497f-4424-a2dd-200118f5ab66.png"), "what did you expect would happen with that outfill method" 431 | for frame in tqdm(list((framesets_path / frameset).iterdir())): 432 | if frame.suffix in [".png", ".jpg"]: 433 | with Image.open(frame) as img: 434 | img = process_image(img, should_resize, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters, None) 435 | save_path = Path(output_dir) 436 | os.makedirs(str(save_path.resolve()), exist_ok=True) 437 | img.save(Path(output_dir) / frame.name) 438 | return gr.update(), f'Processed images saved to "{output_dir}"!' 439 | bulk_process_button.click(fn=bulk_process_button_click, inputs=[frameset_dropdown, resize_checkbox, output_dir, outfill_setting, outfill_color, outfill_border_blur, outfill_n_clusters], outputs=[crop_preview, log_output]) 440 | 441 | def outfill_setting_change(outfill_setting): 442 | outfill_outputs = [ 443 | "outfill_setting_options", 444 | "color_container", 445 | "border_blur_container", 446 | "n_clusters_container", 447 | "original_image_outfill_setting_container" 448 | ] 449 | visibility_pairs = { 450 | "Solid color": [ 451 | "outfill_setting_options", 452 | "color_container" 453 | ], 454 | "Blurred & stretched overlay" : [ 455 | "outfill_setting_options", 456 | "border_blur_container" 457 | ], 458 | "Dominant image color": [ 459 | "outfill_setting_options", 460 | "n_clusters_container" 461 | ], 462 | "Stretch pixels at border": [ 463 | "outfill_setting_options", 464 | "border_blur_container" 465 | ], 466 | "Reflect image around border": [ 467 | "outfill_setting_options", 468 | "border_blur_container" 469 | ], 470 | "Reuse original image": [ 471 | "outfill_setting_options", 472 | "border_blur_container", 473 | "original_image_outfill_setting_container" 474 | ] 475 | } 476 | return [gr.update(visible=(outfill_setting in visibility_pairs and o in visibility_pairs[outfill_setting])) for o in outfill_outputs] 477 | outfill_setting.change(fn=outfill_setting_change, inputs=[outfill_setting], outputs=[outfill_setting_options, color_container, border_blur_container, n_clusters_container, original_image_outfill_setting_container]) 478 | 479 | reset_aspect_ratio_button.click(fn=None, _js="resetAspectRatio", inputs=[], outputs=[]) 480 | 481 | return (training_picker, "Training Picker", "training_picker"), 482 | 483 | def on_ui_settings(): 484 | picker_path = Path(paths.script_path) / "training-picker" 485 | section = ('training-picker', "Training Picker") 486 | opts.add_option("training_picker_fixed_size", OptionInfo(512, "Fixed size to resize images to", section=section)) 487 | opts.add_option("training_picker_videos_path", OptionInfo(str(picker_path / "videos"), "Path to read videos from", section=section)) 488 | opts.add_option("training_picker_framesets_path", OptionInfo(str(picker_path / "extracted-frames"), "Path to store extracted frame sets in", section=section)) 489 | opts.add_option("training_picker_default_output_path", OptionInfo(str(picker_path / "cropped-frames"), "Default cropped image output directory", section=section)) 490 | 491 | script_callbacks.on_ui_settings(on_ui_settings) 492 | script_callbacks.on_ui_tabs(on_ui_tabs) 493 | --------------------------------------------------------------------------------