feat: add support for Tensorflow dataset - #729
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a TensorFlow-facing data-loading path so ApertureDB query results can be consumed as a tf.data.Dataset, intended to mirror the existing PyTorch dataset wrapper and address issue #2.
Changes:
- Introduces
ApertureDBTensorFlowDatasetto stream batchedFindImageresults into atf.data.Datasetviafrom_generator. - Adds a TensorFlow connector test that validates end-to-end iteration and batching/prefetch behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
aperturedb/TensorFlowDataset.py |
New TensorFlow dataset wrapper that fetches images in server-side batches and exposes them as a tf.data.Dataset with inferred label dtype. |
test/test_tf_connector.py |
New integration-style test that iterates through the TensorFlow dataset and validates total item counts, including a batched/prefetched pipeline. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+20
to
+22
| if tf.shape(img)[0] < 0: | ||
| logger.error("Empty image?") | ||
| assert True == False |
Comment on lines
+2
to
+4
| import os | ||
| import logging | ||
| from typing import Union |
- Remove unused imports - Fix tensor shape check condition - Add test for dynamic label dtype
|
|
||
| if self.label_prop: | ||
| entities = r[self.find_image_idx]["FindImage"]["entities"] | ||
| self.batch_labels = [l[self.label_prop] for l in entities] |
| if time_taken != 0: | ||
| logger.info(f"Throughput (imgs/s): {expected_length / time_taken}") | ||
|
|
||
| def test_nativeContraints(self, db, utils, images): |
Comment on lines
+57
to
+85
| def test_datasetWithMultiprocessing(self, db, utils, images): | ||
| len_limit = utils.count_images() | ||
| # This is a hack against a bug in batch API. | ||
| dim = 224 if isinstance(db, ConnectorRest) else 225 | ||
| query = [{ | ||
| "FindImage": { | ||
| "constraints": { | ||
| "age": [">=", 0] | ||
| }, | ||
| "operations": [ | ||
| { | ||
| "type": "resize", | ||
| "width": dim, | ||
| "height": dim | ||
| } | ||
| ], | ||
| "results": { | ||
| "list": ["license"], | ||
| "limit": len_limit | ||
| } | ||
| } | ||
| }] | ||
|
|
||
| dataset_wrapper = ApertureDBTensorFlowDataset( | ||
| db, query, label_prop="license") | ||
| dataset = dataset_wrapper.get_dataset() | ||
| # Test batched dataset | ||
| batched_dataset = dataset.batch(10).prefetch(tf.data.AUTOTUNE) | ||
|
|
Comment on lines
+55
to
+61
| self.query[self.find_image_idx]["FindImage"]["blobs"] = True | ||
|
|
||
| try: | ||
| _, r, b = execute_query( | ||
| client=self.client, query=self.query, blobs=[]) | ||
| batch = r[self.find_image_idx]["FindImage"]["batch"] | ||
| self.total_elements = batch["total_elements"] |
| start = time.time() | ||
| count = 0 | ||
| for imgs, labels in batched_dataset: | ||
| count += imgs.shape[0] |
luisremis
approved these changes
May 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds data loading support from Tensorflow as a replacement for load/preprocess/prepare for training, completing issue #2. This implements an
ApertureDBTensorFlowDatasetmirroring the design ofApertureDBDatasetfor PyTorch, allowing users to easily load image and label batches from ApertureDB directly into atf.data.Datasetwith dynamic label type inference.Verification
Created a unit test
test_tf_connector.pythat creates anApertureDBTensorFlowDatasetand validates it retrieves the correct number of images with both native constraints and simulated multiprocessing environments. Verified that syntax validation and data generation match the dataset structure.Fixes #2