bug fix.

app.py

@@ -51,11 +51,8 @@ def load_mask_ui(input_folder="example_tmp",load_edit = False):
 def load_image_ui(load_edit, input_folder="example_tmp"):
     #try:
     if 1:
-        for img_path in Path(input_folder).iterdir():
-            if img_path.name in ["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')
+        image, 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)
         print("!!", len(mask_np_list))  
         max_val = len(mask_np_list)-1

segment.py

@@ -46,16 +46,16 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
     handles = []
     label_list = []
 
-    mask_list = []
+    mask_np_list = []
     
     if not noseg:
         if torch.min(segmentation) == 0: 
             mask = segmentation==0
-            mask = mask.cpu().detach()   # [512,512]   bool
+            mask = mask.cpu().detach().numpy()   # [512,512]   bool
             segment_label = "rest"
             color = viridis(0)
             label = f"{segment_label}-{0}"
-            mask_list.append(mask)  
+            mask_np_list.append(mask)  
             handles.append(mpatches.Patch(color=color, label=label))
             label_list.append(label)
        
@@ -64,8 +64,8 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             mask = segmentation==segment_id
             if torch.min(segmentation) != 0: 
                 segment_id -= 1
-            mask = mask.cpu().detach()   # [512,512] bool
-            mask_list.append(mask) 
+            mask = mask.cpu().detach().numpy()   # [512,512] bool
+            mask_np_list.append(mask) 
             segment_label = model.config.id2label[segment['label_id']]
             instances_counter[segment['label_id']] += 1
 
@@ -75,9 +75,9 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             handles.append(mpatches.Patch(color=color, label=label))
             label_list.append(label)
     else:
-        mask = torch.from_numpy(np.full(segmentation.shape, True))
+        mask = np.full(segmentation.shape, True)
         segment_label = "all"
-        mask_list.append(mask) 
+        mask_np_list.append(mask) 
         color = viridis(0)
         label = f"{segment_label}-{0}"
         handles.append(mpatches.Patch(color=color, label=label))
@@ -89,7 +89,7 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
     ax.legend(handles=handles)
     plt.savefig(os.path.join(save_folder, 'seg_init.png'), dpi=500 )
     print("; ".join(label_list))
-    return mask_list,label_list
+    return mask_np_list,label_list
 
 
 
@@ -110,7 +110,7 @@ def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
     image =Image.fromarray(image)
     image = image.resize((size, size))
     os.makedirs(name, exist_ok=True)
-    image.save(os.path.join(name,"img_{}.png".format(size)))
+    #image.save(os.path.join(name,"img_{}.png".format(size)))
     inputs = processor(image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
@@ -121,4 +121,4 @@ def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
     mask_list,label_list = draw_panoptic_segmentation(**panoptic_segmentation, save_folder = save_folder, noseg = noseg, model = model)
     print("Finish segment")
     #block_flag += 1
-    return  mask_list,label_list#, gr.Button.update("1.2 Load edited masks",visible = True), gr.Checkbox.update(label = "Show Segmentation",visible = True) 
+    return  image,mask_list,label_list#, gr.Button.update("1.2 Load edited masks",visible = True), gr.Checkbox.update(label = "Show Segmentation",visible = True)