add checks for sample integrity

Orion-AI-Lab · Jun 21, 2023 · 4da286f · 4da286f
1 parent a925e5d
commit 4da286f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/dataset/Dataset.py b/dataset/Dataset.py
@@ -13,6 +13,7 @@
 import einops
 import torchvision
 import albumentations as A
+
 np.random.seed(999)
 torch.manual_seed(999)
 random.seed(999)
@@ -257,13 +258,15 @@ def __init__(self, config, mode="train"):
                 label = 0
             else:
                 label = 1
-
+            sample_insar_path = annotation_utils.get_insar_path(
+                    annotation_path=annotation_path + annotation_file,
+                    root_path=self.config["data_path"])
+            sample_cc_path = sample_insar_path[:-8] + 'cc.png'
+            if not os.path.isfile(sample_cc_path) or not os.path.isfile(sample_insar_path):
+                continue
             sample_dict = {
                 "frameID": annotation["frameID"],
-                "insar_path": annotation_utils.get_insar_path(
-                    annotation_path=annotation_path + annotation_file,
-                    root_path=self.config["data_path"],
-                ),
+                "insar_path":sample_insar_path,
                 "label": annotation,
             }
             self.interferograms.append(sample_dict)

diff --git a/utilities/utils.py b/utilities/utils.py
@@ -194,7 +194,6 @@ def load_checkpoint(model, optimizer, args):
 def extract_state_dict_from_ddp_checkpoint(checkpoint_path):
     print("=> loading checkpoint '{}'".format(checkpoint_path))
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
-    print(checkpoint.keys())
     encoder_state_dict = {}
     for key in list(checkpoint["state_dict"].keys()):
         checkpoint["state_dict"][key.replace("module.", "")] = checkpoint[