fix: tf dataset get item

Author: PJEstrada

sdk/diffgram/core/directory.py

@@ -79,7 +79,7 @@ def __init__(self, client, file_id_list_sliced = None):
 		self.client = client
 		self.id = None
 		self.file_list_metadata = {}
-		
+
 		if file_id_list_sliced is None:
 			self.file_id_list = self.all_file_ids()
 		else:
@@ -145,8 +145,7 @@ def to_tensorflow(self):
 			project = self.client,
 			diffgram_file_id_list = file_id_list
 		)
-		tf_dataset = diffgram_tensorflow_dataset.get_dataset_obj()
-		return tf_dataset
+		return diffgram_tensorflow_dataset
 
 	def new(self, name: str):
 		"""

sdk/diffgram/core/sliced_directory.py

@@ -46,5 +46,4 @@ def to_tensorflow(self):
             project = self.client,
             diffgram_file_id_list = file_id_list
         )
-        tf_dataset = diffgram_tensorflow_dataset.get_dataset_obj()
-        return tf_dataset
+        return diffgram_tensorflow_dataset

sdk/diffgram/tensorflow_diffgram/diffgram_tensorflow_dataset.py

@@ -1,12 +1,14 @@
 from diffgram.core.diffgram_dataset_iterator import DiffgramDatasetIterator
 import os
+
 try:
     import tensorflow as tf  # type: ignore
 except ModuleNotFoundError:
     raise ModuleNotFoundError(
         "'tensorflow' module should be installed to convert the Dataset into tensorflow format"
     )
 
+
 class DiffgramTensorflowDataset(DiffgramDatasetIterator):
 
     def __init__(self, project, diffgram_file_id_list):
@@ -47,17 +49,13 @@ def __validate_file_ids(self):
             raise Exception(
                 'Some file IDs do not belong to the project. Please provide only files from the same project.')
 
-    def __iter__(self):
-        self.current_file_index = 0
-        return self
-
-    def get_next_elm(self):
-        yield self.__next__()
+    def __getitem__(self, idx):
+        tf_example = self.get_tf_train_example(idx)
+        return tf_example
 
-    def __next__(self):
-        file_id = self.diffgram_file_id_list[self.current_file_index]
+    def get_tf_train_example(self, idx):
+        file_id = self.diffgram_file_id_list[idx]
         diffgram_file = self.project.file.get_by_id(file_id, with_instances = True)
-        print('AAA', diffgram_file.id)
         image = self.get_image_data(diffgram_file)
         instance_data = self.get_file_instances(diffgram_file)
         filename, file_extension = os.path.splitext(instance_data['diffgram_file'].image['original_filename'])
@@ -77,8 +75,12 @@ def __next__(self):
             'image/object/class/label': self.int64_list_feature(instance_data['label_id_list']),
         }
         tf_example = tf.train.Example(features = tf.train.Features(feature = tf_example_dict))
+        return tf_example
+
+    def __next__(self):
+        tf_example = self.get_tf_train_example(self.current_file_index)
         self.current_file_index += 1
         return tf_example
 
-    def get_dataset_obj(self):
-        return tf.data.Dataset.from_generator(self.get_next_elm, output_signature = tf.TensorSpec(shape=(1,)))
+    # def get_dataset_obj(self):
+    #     return tf.data.Dataset.from_generator(self.get_next_elm, output_signature = tf.TensorSpec(shape = (1,)))