second

.gitignore

@@ -4,6 +4,7 @@ example1_example2_512/
 example1_example2_1024/
 example1/
 old/
+example_tmp/
 
 out_active.png
 out_mask.png

app copy.py

@@ -0,0 +1,349 @@
+
+import os
+import copy
+from PIL import Image
+import matplotlib 
+import numpy as np
+import gradio as gr
+from utils import load_mask, load_mask_edit
+from utils_mask import process_mask_to_follow_priority, mask_union, visualize_mask_list_clean
+from pathlib import Path
+import subprocess
+from PIL import Image
+
+LENGTH=512 #length of the square area displaying/editing images
+TRANSPARENCY = 150 # transparency of the mask in display
+
+def add_mask(mask_np_list_updated, mask_label_list):
+    mask_new = np.zeros_like(mask_np_list_updated[0])
+    mask_np_list_updated.append(mask_new)
+    mask_label_list.append("new")
+    return mask_np_list_updated, mask_label_list
+
+def create_segmentation(mask_np_list):
+    viridis = matplotlib.pyplot.get_cmap(name = 'viridis', lut = len(mask_np_list))
+    segmentation = 0
+    for i, m  in enumerate(mask_np_list):
+        color = matplotlib.colors.to_rgb(viridis(i))
+        color_mat = np.ones_like(m)                                                                            
+        color_mat = np.stack([color_mat*color[0], color_mat*color[1],color_mat*color[2] ], axis = 2)
+        color_mat = color_mat * m[:,:,np.newaxis]
+        segmentation += color_mat
+    segmentation = Image.fromarray(np.uint8(segmentation*255))
+    return segmentation
+
+def load_mask_ui(input_folder,load_edit = False):
+    if not load_edit:
+        mask_list, mask_label_list = load_mask(input_folder)
+    else:
+        mask_list, mask_label_list = load_mask_edit(input_folder) 
+        
+    mask_np_list = []
+    for  m  in mask_list:
+        mask_np_list. append( m.cpu().numpy())
+
+    return mask_np_list, mask_label_list
+
+def load_image_ui(input_folder, load_edit):
+    try:
+        for img_path in Path(input_folder).iterdir():
+            if img_path.name in ["img.png", "img_1024.png", "img_512.png"]:
+                image = Image.open(img_path)
+        mask_np_list, mask_label_list = load_mask_ui(input_folder, load_edit = load_edit)
+        image = image.convert('RGB')
+        segmentation = create_segmentation(mask_np_list)  
+        return image, segmentation, mask_np_list, mask_label_list, image
+    except:
+        print("Image folder invalid: The folder should contain image.png")
+        return None, None, None, None, None
+
+def run_segmentation(input_folder):
+    subprocess.run(["python", "segment.py" , "--name={}".format(input_folder)])
+    return 
+
+
+
+def run_edit_text(
+        input_folder,
+        num_tokens,
+        num_sampling_steps,
+        strength,
+        edge_thickness,
+        tgt_prompt,
+        tgt_idx,
+        guidance_scale
+    ):
+    subprocess.run(["python", 
+                    "main.py" ,
+                    "--text",
+                    "--name={}".format(input_folder),
+                    "--dpm={}".format("sd"),
+                    "--resolution={}".format(512),
+                    "--load_trained",
+                    "--num_tokens={}".format(num_tokens),
+                    "--seed={}".format(2024),
+                    "--guidance_scale={}".format(guidance_scale),
+                    "--num_sampling_step={}".format(num_sampling_steps),
+                    "--strength={}".format(strength),
+                    "--edge_thickness={}".format(edge_thickness),
+                    "--num_imgs={}".format(2),
+                    "--tgt_prompt={}".format(tgt_prompt) ,
+                    "--tgt_index={}".format(tgt_idx)            
+    ])
+    
+    return Image.open(os.path.join(input_folder, "text", "out_text_0.png"))
+
+
+def run_optimization(
+        input_folder,
+        num_tokens,
+        embedding_learning_rate, 
+        max_emb_train_steps, 
+        diffusion_model_learning_rate, 
+        max_diffusion_train_steps,
+        train_batch_size,
+        gradient_accumulation_steps
+    ):
+    subprocess.run(["python", 
+                    "main.py" ,
+                    "--name={}".format(input_folder),
+                    "--dpm={}".format("sd"),
+                    "--resolution={}".format(512),
+                    "--num_tokens={}".format(num_tokens),
+                    "--embedding_learning_rate={}".format(embedding_learning_rate),
+                    "--diffusion_model_learning_rate={}".format(diffusion_model_learning_rate),
+                    "--max_emb_train_steps={}".format(max_emb_train_steps),
+                    "--max_diffusion_train_steps={}".format(max_diffusion_train_steps),
+                    "--train_batch_size={}".format(train_batch_size),
+                    "--gradient_accumulation_steps={}".format(gradient_accumulation_steps)
+                    
+    ])
+    return 
+
+
+def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
+    backimg_solid_np =  np.array(backimg)
+    bimg = backimg.copy()
+    fimg = foreimg.copy()
+    fimg.putalpha(transparency)
+    bimg.paste(fimg, (0,0), fimg)
+
+    bimg_np = np.array(bimg)
+    mask_np = mask_np[:,:,np.newaxis]
+    try:
+        new_img_np = bimg_np*mask_np + (1-mask_np)* backimg_solid_np
+        return Image.fromarray(new_img_np) 
+    except:
+        import pdb; pdb.set_trace()
+
+def show_segmentation(image, segmentation, flag):
+    if flag is False:
+        flag = True
+        mask_np = np.ones([image.size[0],image.size[1]]).astype(np.uint8)
+        image_edit = transparent_paste_with_mask(image, segmentation, mask_np ,transparency = TRANSPARENCY)
+        return image_edit, flag
+    else:
+        flag = False
+        return image,flag
+
+def edit_mask_add(canvas,  image, idx, mask_np_list):
+    mask_sel = mask_np_list[idx]
+    mask_new = np.uint8(canvas["mask"][:, :, 0]/ 255.)
+    mask_np_list_updated = []
+    for midx, m  in enumerate(mask_np_list):
+        if midx == idx:
+            mask_np_list_updated.append(mask_union(mask_sel, mask_new))
+        else:
+            mask_np_list_updated.append(m)
+    
+    priority_list = [0 for _ in range(len(mask_np_list_updated))]
+    priority_list[idx] = 1
+    mask_np_list_updated = process_mask_to_follow_priority(mask_np_list_updated, priority_list)
+    mask_ones = np.ones([mask_sel.shape[0], mask_sel.shape[1]]).astype(np.uint8)
+    segmentation = create_segmentation(mask_np_list_updated)
+    image_edit = transparent_paste_with_mask(image, segmentation, mask_ones ,transparency = TRANSPARENCY)
+    return mask_np_list_updated, image_edit
+
+def slider_release(index, image,  mask_np_list_updated, mask_label_list):
+    if index > len(mask_np_list_updated):
+        return image, "out of range"
+    else:
+        mask_np = mask_np_list_updated[index]
+        mask_label = mask_label_list[index]
+        segmentation = create_segmentation(mask_np_list_updated)
+        new_image = transparent_paste_with_mask(image, segmentation, mask_np, transparency = TRANSPARENCY)
+    return new_image, mask_label
+
+def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder):
+    try: 
+        assert np.all(sum(mask_np_list_updated)==1)
+    except:
+        print("please check mask")
+        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
+        import pdb; pdb.set_trace()
+        
+    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
+        # np.save(os.path.join(input_folder, "maskEDIT{}_{}.npy".format(midx, mask_label)),mask )
+        np.save(os.path.join(input_folder, "mask{}_{}.npy".format(midx, mask_label)),mask )
+    savepath = os.path.join(input_folder, "seg_current.png")
+    visualize_mask_list_clean(mask_np_list_updated, savepath)
+    
+def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder):
+    try: 
+        assert np.all(sum(mask_np_list_updated)==1)
+    except:
+        print("please check mask")
+        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
+        import pdb; pdb.set_trace()
+    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
+        np.save(os.path.join(input_folder, "maskEdited{}_{}.npy".format(midx, mask_label)), mask)
+    savepath = os.path.join(input_folder, "seg_edited.png")
+    visualize_mask_list_clean(mask_np_list_updated, savepath)
+  
+with gr.Blocks() as demo:
+    image = gr.State() # store mask
+    image_loaded = gr.State()
+    segmentation    = gr.State()
+
+    mask_np_list    = gr.State([])
+    mask_label_list = gr.State([])
+    mask_np_list_updated = gr.State([])
+    true    = gr.State(True)
+    false    = gr.State(False)
+
+
+    with gr.Row():
+        gr.Markdown("""# D-Edit""")
+
+    with gr.Tab(label="1 Edit mask"):
+        with gr.Row():
+            with gr.Column():
+                canvas = gr.Image(value = None, type="numpy",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+                input_folder = gr.Textbox(value="example1", label="input folder", interactive= True, )
+                
+                segment_button  = gr.Button("1.1 Run segmentation")
+                segment_button.click(run_segmentation, 
+                        [input_folder] , 
+                        [] )
+             
+                
+                text_button  = gr.Button("1.2 Load original masks")
+                text_button.click(load_image_ui, 
+                        [input_folder, false] , 
+                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
+            
+                load_edit_button = gr.Button("1.2 Load edited masks")    
+                load_edit_button.click(load_image_ui, 
+                        [input_folder, true] , 
+                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
+                
+                show_segment = gr.Checkbox(label = "Show Segmentation")
+       
+                flag = gr.State(False)
+                show_segment.select(show_segmentation,
+                                    [image_loaded, segmentation, flag], 
+                                    [canvas, flag])
+       
+            mask_np_list_updated = copy.deepcopy(mask_np_list)
+
+            with gr.Column():
+                gr.Markdown("""<p style="text-align: center; font-size: 20px">Draw Mask</p>""")
+                slider =  gr.Slider(0, 20, step=1,  interactive=True)
+                label = gr.Textbox()
+                slider.release(slider_release, 
+                        inputs = [slider, image_loaded,   mask_np_list_updated, mask_label_list], 
+                        outputs= [canvas, label]
+                    )
+                add_button  = gr.Button("Add")
+                add_button.click( edit_mask_add, 
+                        [canvas, image_loaded, slider, mask_np_list_updated] , 
+                        [mask_np_list_updated, canvas]
+                    )
+
+                save_button2  = gr.Button("Set and Save as edited masks")
+                save_button2.click( save_as_edit_mask, 
+                        [mask_np_list_updated,  mask_label_list, input_folder] , 
+                        [] )  
+                
+                save_button  = gr.Button("Set and Save as original masks")
+                save_button.click( save_as_orig_mask, 
+                        [mask_np_list_updated,  mask_label_list, input_folder] , 
+                        [] )  
+                
+                back_button  = gr.Button("Back to current seg")
+                back_button.click( load_mask_ui, 
+                                [input_folder] , 
+                                [ mask_np_list_updated,mask_label_list] )
+
+                add_mask_button = gr.Button("Add new empty mask")    
+                add_mask_button.click(add_mask, 
+                        [mask_np_list_updated, mask_label_list] , 
+                        [mask_np_list_updated, mask_label_list] )
+                
+    with gr.Tab(label="2 Optimization"):
+        with gr.Row():
+            with gr.Column():
+                canvas_opt = gr.Image(value = canvas.value, type="pil",  label="Loaded Image", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+            
+            with gr.Column():
+                gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
+                num_tokens = gr.Textbox(value="5", label="num tokens to represent each object", interactive= True)
+                embedding_learning_rate = gr.Textbox(value="1e-4", label="Embedding optimization: Learning rate", interactive= True )
+                max_emb_train_steps =  gr.Textbox(value="500", label="embedding optimization: Training steps", interactive= True )
+                
+                diffusion_model_learning_rate = gr.Textbox(value="5e-5", label="UNet Optimization: Learning rate", interactive= True )
+                max_diffusion_train_steps = gr.Textbox(value="500", label="UNet Optimization: Learning rate: Training steps", interactive= True )
+                
+                train_batch_size = gr.Textbox(value="5", label="Batch size", interactive= True )
+                gradient_accumulation_steps=gr.Textbox(value="5", label="Gradient accumulation", interactive= True )
+                
+                add_button  = gr.Button("Run optimization")
+                add_button.click(run_optimization, 
+                        inputs = [
+                            input_folder,
+                            num_tokens,
+                            embedding_learning_rate, 
+                            max_emb_train_steps, 
+                            diffusion_model_learning_rate, 
+                            max_diffusion_train_steps,
+                            train_batch_size,gradient_accumulation_steps
+                        ], 
+                        outputs = []
+                )
+
+
+    with gr.Tab(label="3 Editing"):
+        with gr.Tab(label="3.1 Text-based editing"):
+            canvas_text_edit =  gr.State() # store mask
+            with gr.Row():
+                with gr.Column():
+                    canvas_text_edit = gr.Image(value = None, label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
+                    # canvas_text_edit = gr.Gallery(label = "Edited results")
+                
+                with gr.Column():
+                    gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting (SD)</p>""")
+                    
+                    tgt_prompt =  gr.Textbox(value="Dog", label="Editing: Text prompt", interactive= True )
+                    tgt_idx = gr.Textbox(value="0", label="Editing: Object index", interactive= True )
+                    guidance_scale = gr.Textbox(value="6", label="Editing: CFG guidance scale", interactive= True )
+                    num_sampling_steps = gr.Textbox(value="50", label="Editing: Sampling steps", interactive= True )
+                    edge_thickness = gr.Textbox(value="10", label="Editing: Edge thickness", interactive= True )
+                    strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
+                    
+                    add_button  = gr.Button("Run Editing")
+                    add_button.click(run_edit_text, 
+                        inputs = [
+                            input_folder,
+                            num_tokens,
+                            num_sampling_steps,
+                            strength,
+                            edge_thickness,
+                            tgt_prompt,
+                            tgt_idx,
+                            guidance_scale
+                        ], 
+                        outputs = [canvas_text_edit]
+                    )
+
+
+demo.queue().launch(share=True, debug=True)

app.py

@@ -57,12 +57,6 @@ def load_image_ui(input_folder, load_edit):
         print("Image folder invalid: The folder should contain image.png")
         return None, None, None, None, None
 
-def run_segmentation(input_folder):
-    subprocess.run(["python", "segment.py" , "--name={}".format(input_folder)])
-    return 
-
-
-
 def run_edit_text(
         input_folder,
         num_tokens,
@@ -200,6 +194,8 @@ def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder):
     savepath = os.path.join(input_folder, "seg_edited.png")
     visualize_mask_list_clean(mask_np_list_updated, savepath)
   
+  
+from segment import run_segmentation
 with gr.Blocks() as demo:
     image = gr.State() # store mask
     image_loaded = gr.State()
@@ -211,22 +207,20 @@ with gr.Blocks() as demo:
     true    = gr.State(True)
     false    = gr.State(False)
 
-
     with gr.Row():
         gr.Markdown("""# D-Edit""")
 
     with gr.Tab(label="1 Edit mask"):
         with gr.Row():
             with gr.Column():
-                canvas = gr.Image(value = None, type="numpy",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+                canvas = gr.Image(value = None, type="pil",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
                 input_folder = gr.Textbox(value="example1", label="input folder", interactive= True, )
                 
                 segment_button  = gr.Button("1.1 Run segmentation")
                 segment_button.click(run_segmentation, 
-                        [input_folder] , 
+                        [canvas] , 
                         [] )
-             
-                
+
                 text_button  = gr.Button("1.2 Load original masks")
                 text_button.click(load_image_ui, 
                         [input_folder, false] , 
@@ -280,70 +274,70 @@ with gr.Blocks() as demo:
                         [mask_np_list_updated, mask_label_list] , 
                         [mask_np_list_updated, mask_label_list] )
                 
-    with gr.Tab(label="2 Optimization"):
-        with gr.Row():
-            with gr.Column():
-                canvas_opt = gr.Image(value = canvas.value, type="pil",  label="Loaded Image", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+    # with gr.Tab(label="2 Optimization"):
+    #     with gr.Row():
+    #         with gr.Column():
+    #             canvas_opt = gr.Image(value = canvas.value, type="pil",  label="Loaded Image", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
             
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
-                num_tokens = gr.Textbox(value="5", label="num tokens to represent each object", interactive= True)
-                embedding_learning_rate = gr.Textbox(value="1e-4", label="Embedding optimization: Learning rate", interactive= True )
-                max_emb_train_steps =  gr.Textbox(value="500", label="embedding optimization: Training steps", interactive= True )
+    #         with gr.Column():
+    #             gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
+    #             num_tokens = gr.Textbox(value="5", label="num tokens to represent each object", interactive= True)
+    #             embedding_learning_rate = gr.Textbox(value="1e-4", label="Embedding optimization: Learning rate", interactive= True )
+    #             max_emb_train_steps =  gr.Textbox(value="500", label="embedding optimization: Training steps", interactive= True )
                 
-                diffusion_model_learning_rate = gr.Textbox(value="5e-5", label="UNet Optimization: Learning rate", interactive= True )
-                max_diffusion_train_steps = gr.Textbox(value="500", label="UNet Optimization: Learning rate: Training steps", interactive= True )
+    #             diffusion_model_learning_rate = gr.Textbox(value="5e-5", label="UNet Optimization: Learning rate", interactive= True )
+    #             max_diffusion_train_steps = gr.Textbox(value="500", label="UNet Optimization: Learning rate: Training steps", interactive= True )
                 
-                train_batch_size = gr.Textbox(value="5", label="Batch size", interactive= True )
-                gradient_accumulation_steps=gr.Textbox(value="5", label="Gradient accumulation", interactive= True )
+    #             train_batch_size = gr.Textbox(value="5", label="Batch size", interactive= True )
+    #             gradient_accumulation_steps=gr.Textbox(value="5", label="Gradient accumulation", interactive= True )
                 
-                add_button  = gr.Button("Run optimization")
-                add_button.click(run_optimization, 
-                        inputs = [
-                            input_folder,
-                            num_tokens,
-                            embedding_learning_rate, 
-                            max_emb_train_steps, 
-                            diffusion_model_learning_rate, 
-                            max_diffusion_train_steps,
-                            train_batch_size,gradient_accumulation_steps
-                        ], 
-                        outputs = []
-                )
+    #             add_button  = gr.Button("Run optimization")
+    #             add_button.click(run_optimization, 
+    #                     inputs = [
+    #                         input_folder,
+    #                         num_tokens,
+    #                         embedding_learning_rate, 
+    #                         max_emb_train_steps, 
+    #                         diffusion_model_learning_rate, 
+    #                         max_diffusion_train_steps,
+    #                         train_batch_size,gradient_accumulation_steps
+    #                     ], 
+    #                     outputs = []
+    #             )
 
 
-    with gr.Tab(label="3 Editing"):
-        with gr.Tab(label="3.1 Text-based editing"):
-            canvas_text_edit =  gr.State() # store mask
-            with gr.Row():
-                with gr.Column():
-                    canvas_text_edit = gr.Image(value = None, label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
-                    # canvas_text_edit = gr.Gallery(label = "Edited results")
+    # with gr.Tab(label="3 Editing"):
+    #     with gr.Tab(label="3.1 Text-based editing"):
+    #         canvas_text_edit =  gr.State() # store mask
+    #         with gr.Row():
+    #             with gr.Column():
+    #                 canvas_text_edit = gr.Image(value = None, label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
+    #                 # canvas_text_edit = gr.Gallery(label = "Edited results")
                 
-                with gr.Column():
-                    gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting (SD)</p>""")
+    #             with gr.Column():
+    #                 gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting (SD)</p>""")
                     
-                    tgt_prompt =  gr.Textbox(value="Dog", label="Editing: Text prompt", interactive= True )
-                    tgt_idx = gr.Textbox(value="0", label="Editing: Object index", interactive= True )
-                    guidance_scale = gr.Textbox(value="6", label="Editing: CFG guidance scale", interactive= True )
-                    num_sampling_steps = gr.Textbox(value="50", label="Editing: Sampling steps", interactive= True )
-                    edge_thickness = gr.Textbox(value="10", label="Editing: Edge thickness", interactive= True )
-                    strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
+    #                 tgt_prompt =  gr.Textbox(value="Dog", label="Editing: Text prompt", interactive= True )
+    #                 tgt_idx = gr.Textbox(value="0", label="Editing: Object index", interactive= True )
+    #                 guidance_scale = gr.Textbox(value="6", label="Editing: CFG guidance scale", interactive= True )
+    #                 num_sampling_steps = gr.Textbox(value="50", label="Editing: Sampling steps", interactive= True )
+    #                 edge_thickness = gr.Textbox(value="10", label="Editing: Edge thickness", interactive= True )
+    #                 strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
                     
-                    add_button  = gr.Button("Run Editing")
-                    add_button.click(run_edit_text, 
-                        inputs = [
-                            input_folder,
-                            num_tokens,
-                            num_sampling_steps,
-                            strength,
-                            edge_thickness,
-                            tgt_prompt,
-                            tgt_idx,
-                            guidance_scale
-                        ], 
-                        outputs = [canvas_text_edit]
-                    )
+    #                 add_button  = gr.Button("Run Editing")
+    #                 add_button.click(run_edit_text, 
+    #                     inputs = [
+    #                         input_folder,
+    #                         num_tokens,
+    #                         num_sampling_steps,
+    #                         strength,
+    #                         edge_thickness,
+    #                         tgt_prompt,
+    #                         tgt_idx,
+    #                         guidance_scale
+    #                     ], 
+    #                     outputs = [canvas_text_edit]
+    #                 )
 
 
 demo.queue().launch(share=True, debug=True)

segment.py

@@ -32,7 +32,7 @@ def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):
     image = np.array(Image.fromarray(image).resize((size, size)))
     return image
 
-def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, noseg = False):
+def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, noseg = False, model =None):
     if torch.max(segmentation)==torch.min(segmentation)==-1:
         print("nothing is detected!")
         noseg=True
@@ -88,28 +88,30 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
 
 
 
-parser = argparse.ArgumentParser()
-parser.add_argument("--name",  type=str, default="obama")
-parser.add_argument("--size",  type=int, default=512)
-parser.add_argument("--noseg", default=False, action="store_true" )
-args = parser.parse_args()
-base_folder_path = "."
 
-processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
-model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
-input_folder = os.path.join(base_folder_path, args.name )
-try:
-    image = load_image(os.path.join(input_folder, "img.png" ), size = args.size)
-except:
-    image = load_image(os.path.join(input_folder, "img.jpg" ), size = args.size)
+def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
 
-image =Image.fromarray(image)
-image.save(os.path.join(input_folder,"img_{}.png".format(args.size)))
-inputs = processor(image, return_tensors="pt")
-with torch.no_grad():
-    outputs = model(**inputs)
+    base_folder_path = "."
 
-panoptic_segmentation = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
-save_folder = os.path.join(base_folder_path, args.name)
-os.makedirs(save_folder, exist_ok=True)
-draw_panoptic_segmentation(**panoptic_segmentation, save_folder = save_folder, noseg = args.noseg)
+    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+
+
+    # input_folder = os.path.join(base_folder_path, name )
+    # try:
+    #     image = load_image(os.path.join(input_folder, "img.png" ), size = size)
+    # except:
+    #     image = load_image(os.path.join(input_folder, "img.jpg" ), size = size)
+    # image =Image.fromarray(image)
+    os.makedirs(name, exist_ok=True)
+    image.save(os.path.join(name,"img_{}.png".format(size)))
+    inputs = processor(image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    panoptic_segmentation = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+    save_folder = os.path.join(base_folder_path, name)
+    os.makedirs(save_folder, exist_ok=True)
+    draw_panoptic_segmentation(**panoptic_segmentation, save_folder = save_folder, noseg = noseg, model = model)
+    print("Finish segment")
+    return