fix bug

UTAustin-SwarmLab · Feb 11, 2024 · cef85a4 · cef85a4
1 parent a6a1457
commit cef85a4
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/run_scripts/run_synthetic_tlv_coco.py b/run_scripts/run_synthetic_tlv_coco.py
@@ -24,7 +24,7 @@
         type=str,
         default="",
     )
-    # - - - - - - Image Loader Arguement - - - - - - #
+    # - - - - - - Image Loader Argument - - - - - - #
     parser.add_argument(
         "--coco_image_source",
         type=str,

diff --git a/tlv_dataset/data/tlv_dataset.py b/tlv_dataset/data/tlv_dataset.py
@@ -35,6 +35,16 @@ def __post_init__(self):
             data=self.frames_of_interest
         )
 
+    def sanity_check(self):
+        """Sanity check."""
+        _labels_of_frames = []
+        for label in self.labels_of_frames:
+            if isinstance(label, str):
+                _labels_of_frames.append([label])
+            else:
+                _labels_of_frames.append(label)
+        self.labels_of_frames = _labels_of_frames
+
     def save_frames(
         self, path="/opt/Neuro-Symbolic-Video-Frame-Search/artifacts"
     ) -> None:

diff --git a/tlv_dataset/generator/synthetic_tlv_generator.py b/tlv_dataset/generator/synthetic_tlv_generator.py
@@ -123,6 +123,8 @@ def generate(
                     proposition=ltl_frame.proposition,
                 )
                 self.extract_properties(ltl_frame.ltl_formula)
+                # sanity check
+                ltl_frame.sanity_check()
                 if save_as == "dict":
                     ltl_frame.save_as_dict(
                         save_path=self._save_dir
@@ -192,6 +194,7 @@ def generate_until_time_delta(
                     proposition=ltl_frame.proposition,
                 )
                 self.extract_properties(ltl_frame.ltl_formula)
+                ltl_frame.sanity_check()
                 if save_as == "dict":
                     ltl_frame.save_as_dict(
                         save_path=self._save_dir
@@ -455,9 +458,11 @@ def ltl_function(
                         if x not in temp_frames_of_interest
                     ]
 
-            # 2. G "prop1"
-
         # TODO: Make a false case
+        # for i, label in enumerate(labels_of_frame):
+        #     if isinstance(label, str):
+        #         labels_of_frame[i] = [label]
+
         return TLVDataset(
             ground_truth=True,
             ltl_formula=ltl_formula,