Update main.py

main.py

@@ -12,7 +12,7 @@ from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualiz
 def run_main(
     name="example_tmp",
     name_2=None,
-    mask_list=None, 
+    mask_np_list=None, 
     mask_label_list=None,
     image_gt=None,
     dpm="sd",
@@ -74,7 +74,11 @@ def run_main(
     base_output_folder  = "."
 
     input_folder = os.path.join(base_input_folder, name)
-
+    mask_list = []
+    for mask_np in mask_np_list:
+        mask = torch.from_numpy(mask_np)
+        mask_list.append(mask)
+    
     #mask_list, mask_label_list = load_mask(input_folder)
     assert mask_list[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution)
     #try: