append probablities to the output

demo.py

@@ -31,6 +31,11 @@ def demo(folder_path, output_path=Path("tmp")):
             image = image.to(opt.device).unsqueeze(0)
             outputs = model(image, seg_size=image_size)
             out_map = outputs["ensemble"]["out_map"][0, ...].detach().cpu()
+            pred = outputs["ensemble"]["out_map"].max().item() 
+            if  pred > opt.mask_threshold:
+                print(f"Found manipulation in {image_path.name}")
+            else:
+                print(f"No manipulation found in {image_path.name}")
 
             overlay = draw_segmentation_masks(
                 dsm_image, masks=out_map[0, ...] > opt.mask_threshold
@@ -43,7 +48,8 @@ def demo(folder_path, output_path=Path("tmp")):
                 ],
                 padding=5,
             )
-            save_image(grid_image, (output_path / image_path.name).as_posix())
+            image_name = image_path.stem + f"-{pred:.2f}" + image_path.suffix
+            save_image(grid_image, (output_path / image_name).as_posix())
 
 
 if __name__ == "__main__":

opt.py

@@ -10,8 +10,8 @@ import yaml
 from termcolor import cprint
 
 
-def load_dataset_arguments(opt):
-    if opt.load is None:
+def load_dataset_arguments(cfg_path, opt):
+    if opt.load is None and cfg_path is None:
         return
 
     # exclude parameters assigned in the command
@@ -24,7 +24,10 @@ def load_dataset_arguments(opt):
         arguments = []
 
     # load parameters in the yaml file
-    assert os.path.exists(opt.load)
+    if cfg_path is not None:
+        opt.load = cfg_path
+    else:
+        assert os.path.exists(opt.load)
     with open(opt.load, "r") as f:
         yaml_arguments = yaml.safe_load(f)
     # TODO this should be verified
@@ -33,7 +36,7 @@ def load_dataset_arguments(opt):
             setattr(opt, k, v)
 
 
-def get_opt(additional_parsers: Optional[List] = None):
+def get_opt(cfg_path: Optional[str] = None, additional_parsers: Optional[List] = None):
     parents = [get_arguments_parser()]
     if additional_parsers:
         parents.extend(additional_parsers)
@@ -43,7 +46,7 @@ def get_opt(additional_parsers: Optional[List] = None):
     opt = parser.parse_known_args()[0]
 
     # load dataset argument file
-    load_dataset_arguments(opt)
+    load_dataset_arguments(cfg_path, opt)
 
     # user-defined warnings and assertions
     if opt.decoder.lower() not in ["c1"]:

requirements.txt

@@ -26,3 +26,4 @@ timm==0.9.12
 torch==1.12.1+cu116
 torchvision==0.13.1+cu116
 tqdm==4.64.1
+markupsafe==2.0.1