Update ade20k_main.py

ccivit · Sep 22, 2020 · c31f472 · c31f472
1 parent adcf2fa
commit c31f472
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/experiments/segmentation/ade20k_main.py b/experiments/segmentation/ade20k_main.py
@@ -55,7 +55,7 @@ def strip_dir_structure(image_dir_path,pipeline_path):
 def load_config(config_file):
     with open(config_file, 'r') as stream:
         config = yaml.safe_load(stream)
-        results_file = config[model_name]['results']
+        results_file = config['image_segmentation']['results']
         pipeline = config['general']['pipeline_path']
     output_dir = strip_dir_structure(img_dir,pipeline)
     return config, results_file, pipeline, output_dir
@@ -66,7 +66,7 @@ def load_config(config_file):
 img_dir = os.path.join(mapped_dir,sys.argv[1])
 config_file = os.path.join(mapped_dir,sys.argv[2])
 config, results_file, pipeline, output_dir = load_config(config_file)
-decoder_file = os.path.join(mapped_dir,config[model_name]['decoder'])
+decoder_file = os.path.join(mapped_dir,config['image_segmentation']['decoder'])
 decoder = load_decoder(decoder_file)
 total_imgs = len(os.listdir(img_dir))