Spaces:
Runtime error
Runtime error
File size: 8,401 Bytes
688976a e8bad6c a6933f9 bb21642 688976a 97f1bae 688976a a6933f9 688976a e8bad6c a6933f9 688976a 2e0064d 51f1b25 a6933f9 2e0064d e8bad6c 4d105da 97f1bae e8bad6c bb21642 c89e57a a6933f9 2e0064d a6933f9 2e0064d a6933f9 2e0064d a6933f9 e8bad6c a6933f9 2e0064d a6933f9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import gradio as gr
from gradio_image_prompter import ImagePrompter
from gradio_image_prompter.image_prompter import PromptData
from typing import List, Dict, Optional, Union
import os
import yaml
from modules.sam_inference import SamInference
from modules.model_downloader import DEFAULT_MODEL_TYPE
from modules.paths import (OUTPUT_DIR, OUTPUT_PSD_DIR, SAM2_CONFIGS_DIR, TEMP_DIR)
from modules.utils import open_folder
from modules.constants import (AUTOMATIC_MODE, BOX_PROMPT_MODE)
from modules.video_utils import extract_frames, get_frames_from_dir
class App:
def __init__(self,
args=None):
self.demo = gr.Blocks()
self.args = args
self.sam_inf = SamInference()
self.image_modes = [AUTOMATIC_MODE, BOX_PROMPT_MODE]
self.default_mode = BOX_PROMPT_MODE
default_param_config_path = os.path.join(SAM2_CONFIGS_DIR, "default_hparams.yaml")
with open(default_param_config_path, 'r') as file:
self.hparams = yaml.safe_load(file)
def mask_parameters(self,
hparams: Optional[Dict] = None):
if hparams is None:
hparams = self.hparams["mask_hparams"]
mask_components = [
gr.Number(label="points_per_side ", value=hparams["points_per_side"], interactive=True),
gr.Number(label="points_per_batch ", value=hparams["points_per_batch"], interactive=True),
gr.Slider(label="pred_iou_thresh ", value=hparams["pred_iou_thresh"], minimum=0, maximum=1,
interactive=True),
gr.Slider(label="stability_score_thresh ", value=hparams["stability_score_thresh"], minimum=0,
maximum=1, interactive=True),
gr.Slider(label="stability_score_offset ", value=hparams["stability_score_offset"], minimum=0,
maximum=1),
gr.Number(label="crop_n_layers ", value=hparams["crop_n_layers"]),
gr.Slider(label="box_nms_thresh ", value=hparams["box_nms_thresh"], minimum=0, maximum=1),
gr.Number(label="crop_n_points_downscale_factor ", value=hparams["crop_n_points_downscale_factor"]),
gr.Number(label="min_mask_region_area ", value=hparams["min_mask_region_area"]),
gr.Checkbox(label="use_m2m ", value=hparams["use_m2m"])
]
return mask_components
@staticmethod
def on_mode_change(mode: str):
return [
gr.Image(visible=mode == AUTOMATIC_MODE),
ImagePrompter(visible=mode == BOX_PROMPT_MODE),
gr.Accordion(visible=mode == AUTOMATIC_MODE),
]
def on_video_upload(self, vid_input: str):
output_temp_dir = TEMP_DIR
extract_frames(vid_input=vid_input, output_temp_dir=output_temp_dir)
frames = get_frames_from_dir(vid_dir=output_temp_dir)
# self.sam_inf.init_video_inference_state(output_temp_dir)
return [
ImagePrompter(label="Prompt image with Box & Point", value=frames[0]),
gr.Slider(label="Frame Indexes", value=0, interactive=True, step=1, minimum=0, maximum=(len(frames)-1))
]
@staticmethod
def on_frame_change(frame_idx: int):
temp_dir = TEMP_DIR
frames = get_frames_from_dir(vid_dir=temp_dir)
selected_frame = frames[frame_idx]
return ImagePrompter(elem_id="vid-prompter-index", label=f"Prompt image with Box & Point #{frame_idx}",
value=selected_frame)
@staticmethod
def on_prompt_change(prompt: Dict):
image, points = prompt["image"], prompt["points"]
return gr.Image(label="Preview", value=image)
def launch(self):
_mask_hparams = self.hparams["mask_hparams"]
with self.demo:
with gr.Tabs():
with gr.TabItem("Layer Divider"):
with gr.Row():
with gr.Column(scale=5):
img_input = gr.Image(label="Input image here", visible=self.default_mode == AUTOMATIC_MODE)
img_input_prompter = ImagePrompter(label="Prompt image with Box & Point", type='pil',
visible=self.default_mode == BOX_PROMPT_MODE)
with gr.Column(scale=5):
dd_input_modes = gr.Dropdown(label="Image Input Mode", value=self.default_mode,
choices=self.image_modes)
dd_models = gr.Dropdown(label="Model", value=DEFAULT_MODEL_TYPE,
choices=self.sam_inf.available_models)
with gr.Accordion("Mask Parameters", open=False, visible=self.default_mode == AUTOMATIC_MODE) as acc_mask_hparams:
mask_hparams_component = self.mask_parameters(_mask_hparams)
cb_multimask_output = gr.Checkbox(label="multimask_output", value=_mask_hparams["multimask_output"])
with gr.Row():
btn_generate = gr.Button("GENERATE", variant="primary")
with gr.Row():
gallery_output = gr.Gallery(label="Output images will be shown here")
with gr.Column():
output_file = gr.File(label="Generated psd file", scale=9)
btn_open_folder = gr.Button("📁\nOpen PSD folder", scale=1)
sources = [img_input, img_input_prompter, dd_input_modes]
model_params = [dd_models]
mask_hparams = mask_hparams_component + [cb_multimask_output]
input_params = sources + model_params + mask_hparams
btn_generate.click(fn=self.sam_inf.divide_layer,
inputs=input_params, outputs=[gallery_output, output_file])
btn_open_folder.click(fn=lambda: open_folder(os.path.join(OUTPUT_PSD_DIR)),
inputs=None, outputs=None)
dd_input_modes.change(fn=self.on_mode_change,
inputs=[dd_input_modes],
outputs=[img_input, img_input_prompter, acc_mask_hparams])
with gr.TabItem("Mosaic Filter"):
with gr.Row(equal_height=True):
with gr.Column(scale=2):
vid_input = gr.Video(label="Input Video here", scale=3)
with gr.Column(scale=8):
with gr.Row():
vid_frame_prompter = ImagePrompter(elem_id="vid-prompter",
label="Prompt image with Box & Point ",
interactive=True, scale=5)
img_preview = gr.Image(label="Preview", interactive=False, scale=5)
sld_frame_selector = gr.Slider(label="Frame Index", interactive=False)
with gr.Row():
btn_generate = gr.Button("GENERATE", variant="primary")
with gr.Row():
gallery_output = gr.Gallery(label="Output images will be shown here")
with gr.Column():
output_file = gr.File(label="Generated psd file", scale=9)
btn_open_folder = gr.Button("📁\nOpen PSD folder", scale=1)
vid_input.change(fn=self.on_video_upload,
inputs=[vid_input],
outputs=[vid_frame_prompter, sld_frame_selector])
sld_frame_selector.change(fn=self.on_frame_change,
inputs=[sld_frame_selector],
outputs=[vid_frame_prompter],)
vid_frame_prompter.change(fn=self.on_prompt_change,
inputs=[vid_frame_prompter],
outputs=[img_preview])
self.demo.queue().launch(inbrowser=True)
if __name__ == "__main__":
demo = App()
demo.launch()
|