diff --git a/Dockerfile b/Dockerfile
index c1f0145878adc5f12c230400bb74fe9337b1f20c..d5bcbfb123b5de447a115c99162fc0d4cb1fa538 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -55,3 +55,4 @@ RUN pip install mmcv-full==1.4.0 -f https://download.openmmlab.com/mmcv/dist/cu1
 # Copy user files
 COPY --chown=1001:1001 . ${HOME_DIR}
 #CMD while true; do sleep 1; done
+
diff --git a/predict.py b/predict.py
index e4798c61e3248c3b307e53c24e98c0ba1ea7dcfb..3351f8f47d961ca5db89088af9ddc4322fb3b56a 100644
--- a/predict.py
+++ b/predict.py
@@ -1,6 +1,7 @@
 import argparse
 import json
 import shutil
+from black import out
 import mmcv
 import os
 import os.path as osp
@@ -24,8 +25,8 @@ def parse_args():
     )
     parser.add_argument("--cfg", help="mmdetection config")
     parser.add_argument("--model", help="mmdetection model")
-    parser.add_argument("--img_dir", help="directory with test images")
-    parser.add_argument("--out", help="path to json output")
+    # parser.add_argument("--img_dir", help="directory with test images")
+    # parser.add_argument("--out", help="path to json output")
 
     args = parser.parse_args()
     return args
@@ -2902,7 +2903,7 @@ def predict(cfg_path, model_path, img_dir, out_fp):
     outputs = single_gpu_test(model, data_loader)
 
     dataset.results2json(outputs, "results")
-    shutil.move('results.segm.json', out_fp)
+    shutil.move("results.segm.json", out_fp)
 
     return outputs
 
@@ -2910,8 +2911,14 @@ def predict(cfg_path, model_path, img_dir, out_fp):
 def main():
     args = parse_args()
 
-    create_empty_json(str(args.img_dir))
-    result = predict(args.cfg, args.model, str(args.img_dir), str(args.out))
+    img_dir = os.getenv("AICROWD_TEST_IMAGES_PATH", False)
+    out_fp = os.getenv("AICROWD_PREDICTIONS_OUTPUT_PATH", False)
+
+    print(img_dir)
+    print(out_fp)
+
+    create_empty_json(img_dir)
+    result = predict(args.cfg, args.model, img_dir, out_fp)
 
 
 if __name__ == "__main__":
diff --git a/run.sh b/run.sh
index fa104c7392e4cb37a62a95b3799e908e9fc87bf1..3a6de3e5a8e161ec985c44a6eff85bb5616f9896 100755
--- a/run.sh
+++ b/run.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
-python predict.py --cfg ./configs/solo.py --model ./models/epoch_20.pth --img_dir "$AICROWD_TEST_IMAGES_PATH" --out "$AICROWD_PREDICTIONS_OUTPUT_PATH"
+python predict.py --cfg ./configs/solo.py --model ./models/epoch_20.pth