preprocess done - pos_pairs are created

contrastive_learning/datasets/preprocess.py

@@ -5,6 +5,8 @@
 
 from tqdm import tqdm
 
+from unitree_legged_msgs.msg import HighCmd
+
 # Methods to preprocess data - is used in dataset
 
 # Dumping video to images
@@ -34,18 +36,56 @@ def dump_video_to_images(root : str, video_type='color') -> None:
 
     print(f'dumping finished in {root}')
 
-def create_pos_pairs(root : str, frame_interval : int, video_type='color') -> None:
-    images_folder = os.path.join(root, f'{video_type}_images')
-    image_names = glob.glob(os.path.join(images_folder, 'frame*'))
+def create_pos_pairs(data_dir : str, video_type='color', frame_interval: int = 5) -> None:
+    images_folder = os.path.join(data_dir, f'videos/{video_type}_images')
+    image_names = glob.glob(os.path.join(images_folder, 'frame_*.jpg'))
     image_names = sorted(image_names)
 
-    # TODO: get actions correctly - currently we're not even adding actions there
+    # Get actions as well
+    with open(os.path.join(data_dir, 'commands.pickle'), 'rb') as f:
+        commands = pickle.load(f)
+
+    print('len(commands): {}, len(image_names): {}'.format(len(commands), len(image_names)))
+    assert len(commands) == len(image_names), "Commands and images don't have the same size"
+
+    # NOTE: Sometimes image_names has a larger size by 1 but we can just get the minimum size
+    min_size = min(len(commands), len(image_names))
+
     pos_pairs = [] # NOTE: this is only images and the next frame with the frame_interval
-    for i in range(len(image_names)-frame_interval):
+    i = 0
+    while True:
+        j = i+1 # We will get separate actions
+        action = (commands[i].forwardSpeed, commands[i].rotateSpeed)
+        while j < min_size-1 and j < i+frame_interval and cmds_are_same(commands[i], commands[j]):
+            j += 1
+        action_j = (commands[j].forwardSpeed, commands[j].rotateSpeed)
+        # print(f'commands[i={i}]: {action} - commands[j={j}]: {action_j}')
+
         pos_pairs.append((
             image_names[i],
-            image_names[i+frame_interval]
+            image_names[j],
+            action
         ))
 
-    with open(os.path.join(root, f'{video_type}_pos_pairs.pkl'), 'wb') as f:
-        pickle.dump(pos_pairs, f) # These pos_pairs files are used in dataset
+        if j == min_size-1:
+            break
+        i = j
+
+    with open(os.path.join(data_dir, f'{video_type}_pos_pairs.pkl'), 'wb') as f:
+        pickle.dump(pos_pairs, f) # These pos_pairs files are used in dataset
+
+def cmds_are_same(cmd_a: HighCmd, cmd_b: HighCmd) -> None: # Gets high level commands and compares forwardSpeed and rotateSpeed
+    return cmd_a.forwardSpeed == cmd_b.forwardSpeed and cmd_a.rotateSpeed == cmd_b.rotateSpeed
+
+
+if __name__ == "__main__":
+    data_dir = "/home/irmak/Workspace/DAWGE/src/dawge_planner/data/box_a_1"
+    data_dirs = glob.glob("/home/irmak/Workspace/DAWGE/src/dawge_planner/data/box_a_*")
+    data_dirs = sorted(data_dirs)
+    print('data_dirs: {}'.format(data_dirs))
+    video_type = 'color'
+
+    for data_dir in data_dirs:
+        print(data_dir)
+        create_pos_pairs(data_dir, video_type, frame_interval=8)
+