From cef85a4675d369099b60d52313440845647ba0a0 Mon Sep 17 00:00:00 2001
From: "Minkyu Choi (mc76728)" <minkyu.choi@utexas.edu>
Date: Sun, 11 Feb 2024 00:09:17 -0600
Subject: [PATCH] fix bug

---
 run_scripts/run_synthetic_tlv_coco.py            |  2 +-
 tlv_dataset/data/tlv_dataset.py                  | 10 ++++++++++
 tlv_dataset/generator/synthetic_tlv_generator.py |  9 +++++++--
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/run_scripts/run_synthetic_tlv_coco.py b/run_scripts/run_synthetic_tlv_coco.py
index f130bbf..f93a460 100644
--- 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
index 1dfc38b..4a8d4f3 100644
--- 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
index 268062f..88ddbab 100644
--- 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,