diff --git a/predict.py b/predict.py index afca7e378e8cc8c6c132598f99e9dd6a7d31e114..4270cc7417fa612a84c46b5dc6d230b79388b136 100644 --- a/predict.py +++ b/predict.py @@ -2881,7 +2881,7 @@ def create_empty_json(img_dir, out_fp="./data/test.json"): json.dump(json_file, f, indent=2) -def predict(cfg_path, model_path, img_dir, out_fp): +def predict(cfg_path, model_path, img_dir, out_dir): cfg = Config.fromfile(cfg_path) cfg.model.train_cfg = None cfg.data.test.test_mode = True @@ -2899,6 +2899,7 @@ def predict(cfg_path, model_path, img_dir, out_fp): outputs = single_gpu_test(model, data_loader) dataset.results2json(outputs, "results") + out_fp = os.path.join(out_dir, "results.segm.json") shutil.move("results.segm.json", out_fp) return outputs @@ -2907,19 +2908,17 @@ def predict(cfg_path, model_path, img_dir, out_fp): def main(): args = parse_args() - img_dir = os.getenv("AICROWD_TEST_IMAGES_PATH", False) - out_fp = os.getenv("AICROWD_PREDICTIONS_OUTPUT_PATH", False) + # img_dir = os.getenv("AICROWD_TEST_IMAGES_PATH", False) + # out_fp = os.getenv("AICROWD_PREDICTIONS_OUTPUT_PATH", False) - debug_test_dir = os.getenv("TEST_DATASET_PATH", False) - debug_out_dir = os.getenv("AICROWD_OUTPUTS_DIR", False) - print(os.listdir(debug_test_dir)) - print(os.listdir(debug_out_dir)) + test_dir = os.getenv("TEST_DATASET_PATH", False) + out_dir = os.getenv("AICROWD_OUTPUTS_DIR", False) - print(img_dir) - print(out_fp) + img_dir = os.path.join(test_dir, "images") + print("DEBUG TEST DIR:", len(os.listdir(img_dir))) create_empty_json(img_dir) - result = predict(args.cfg, args.model, img_dir, out_fp) + result = predict(args.cfg, args.model, img_dir, out_dir) if __name__ == "__main__":