Incorporate PR feedback

* Add empty records concept to tub
* Minor updates to conversion script

donkeycar/parts/tub_v2.py

@@ -9,10 +9,13 @@
 
 
 class Tub(object):
-    '''
+    """
     A datastore to store sensor data in a key, value format. \n
     Accepts str, int, float, image_array, image, and array data types.
-    '''
+    """
+    # static variable for empty records
+    empty_record = {'__empty__': True}
+
     def __init__(self, base_path, inputs=[], types=[], metadata=[],
                  max_catalog_len=1000, read_only=False):
         self.base_path = base_path
@@ -28,15 +31,17 @@ def __init__(self, base_path, inputs=[], types=[], metadata=[],
         if not os.path.exists(self.images_base_path):
             os.makedirs(self.images_base_path, exist_ok=True)
 
-    def write_record(self, record):
-        '''
+    def write_record(self, record=None):
+        """
         Can handle various data types including images.
-        '''
+        """
+        if not record:
+            record = self.empty_record
         contents = dict()
         for key, value in record.items():
             if value is None:
                 continue
-            elif not key in self.input_types:
+            elif key not in self.input_types:
                 continue
             else:
                 input_type = self.input_types[key]
@@ -99,9 +104,9 @@ def _image_file_name(cls, index, key, extension='.jpg'):
 
 
 class TubWriter(object):
-    '''
+    """
     A Donkey part, which can write records to the datastore.
-    '''
+    """
     def __init__(self, base_path, inputs=[], types=[], metadata=[],
                  max_catalog_len=1000):
         self.tub = Tub(base_path, inputs, types, metadata, max_catalog_len)

scripts/convert_to_tub_v2.py

@@ -35,28 +35,32 @@ def convert_to_tub_v2(paths, output_path):
 
         record_paths = legacy_tub.gather_records()
         bar = IncrementalBar('Converting', max=len(record_paths))
-        prev_rec_num = None
+        previous_index = None
         for record_path in record_paths:
             try:
                 contents = Path(record_path).read_text()
                 record = json.loads(contents)
                 image_path = record['cam/image_array']
-                rec_num = int(image_path.split('_')[0])
+                current_index = int(image_path.split('_')[0])
                 image_path = os.path.join(legacy_tub.path, image_path)
                 image_data = Image.open(image_path)
                 record['cam/image_array'] = image_data
                 # first record or they are continuous, just append
-                if not prev_rec_num or rec_num == prev_rec_num + 1:
+                if not previous_index or current_index == previous_index + 1:
                     output_tub.write_record(record)
-                    prev_rec_num = rec_num
+                    previous_index = current_index
                 # otherwise fill the gap with dummy records
                 else:
-                    prev_rec_num += 1
-                    while prev_rec_num < rec_num:
+                    # Skipping over previous record here because it has
+                    # already been written.
+                    previous_index += 1
+                    # Adding empty record nodes, and marking them deleted
+                    # until the next valid record.
+                    while previous_index < current_index:
                         idx = output_tub.manifest.current_index
-                        output_tub.write_record({})
+                        output_tub.write_record()
                         output_tub.delete_record(idx)
-                        prev_rec_num += 1
+                        previous_index += 1
                 bar.next()
             except Exception as exception:
                 print(f'Ignoring record path {record_path}\n', exception)