Add initial UI for mosaic filter

app.py

@@ -1,20 +1,22 @@
 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)
+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.app = gr.Blocks()
+        self.demo = gr.Blocks()
         self.args = args
         self.sam_inf = SamInference()
         self.image_modes = [AUTOMATIC_MODE, BOX_PROMPT_MODE]
@@ -52,51 +54,106 @@ class App:
             gr.Accordion(visible=mode == AUTOMATIC_MODE),
         ]
 
-    def launch(self):
-        _mask_hparams = self.hparams["mask_hparams"]
-
-        with self.app:
-            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"])
+    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))
+        ]
 
-            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)
+    @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)
 
-            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
+    @staticmethod
+    def on_prompt_change(prompt: Dict):
+        image, points = prompt["image"], prompt["points"]
+        return gr.Image(label="Preview", value=image)
 
-            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])
+    def launch(self):
+        _mask_hparams = self.hparams["mask_hparams"]
 
-        self.app.queue().launch(inbrowser=True)
+        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__":
-    app = App()
-    app.launch()
+    demo = App()
+    demo.launch()