Fix asqa_dataset_builder for numpy2 by using int64 for sample_id feature.

PiperOrigin-RevId: 666805583

Author: pierrot0

tensorflow_datasets/core/beam_utils_test.py

@@ -51,6 +51,9 @@ def test_read_from_tfds(
             dummy_dataset, split=split, workers_per_shard=workers_per_shard
         )
         | beam.Map(dataset_utils.as_numpy)
+        # Post numpy2, we don't get `{'id': 0}` but
+        # `{'id': np.int64(0)}`
+        | beam.Map(lambda x: {'id': int(x['id'])})
         | beam.io.WriteToText(os.fspath(tmp_path / 'out.txt'))
     )
 

tensorflow_datasets/core/features/bounding_boxes.py

@@ -88,7 +88,10 @@ def encode_example(self, bbox: Union[bb_utils.BBox, np.ndarray]):
 
     # Validate the coordinates
     for coordinate in bbox:
-      if not isinstance(coordinate, float):
+      if not isinstance(
+          coordinate,
+          (float, np.float16, np.float32, np.float64),
+      ):
         raise ValueError(
             'BBox coordinates should be float. Got {}.'.format(bbox)
         )

tensorflow_datasets/core/utils/image_utils.py

@@ -127,15 +127,17 @@ def get_colormap() -> np.ndarray:
   """
   colormap_path = resource_utils.tfds_path() / 'core/utils/colormap.csv'
   with colormap_path.open() as f:
-    return np.array(list(csv.reader(f)), dtype=np.uint8)
+    colormap = np.array(list(csv.reader(f)), dtype=np.uint8)
+  assert colormap.shape == (256, 3)
+  return colormap
 
 
 def apply_colormap(image: np.ndarray) -> np.ndarray:
   """Apply colormap from grayscale (h, w, 1) to colored (h, w, 3) image."""
   image = image.squeeze(axis=-1)  # (h, w, 1) -> (h, w)
   cmap = get_colormap()  # Get the (256, 3) colormap
-  # Normalize uint16 and convert each value to a unique color
-  return cmap[image % len(cmap)]
+  # Convert the image to uint64 first to avoid overflow.
+  return cmap[(image.astype(np.int64) % 256).astype(np.uint8)]
 
 
 # Visualization single image

tensorflow_datasets/datasets/asqa/asqa_dataset_builder.py

@@ -23,7 +23,7 @@
 
 def _features():
   return tfds.features.FeaturesDict({
-      'sample_id': np.int32,
+      'sample_id': np.int64,
       'ambiguous_question': tfds.features.Text(
           doc='Disambiguated question from AmbigQA.'
       ),
@@ -82,6 +82,7 @@ class Builder(tfds.core.GeneratorBasedBuilder):
 
   VERSION = tfds.core.Version('1.0.0')
   RELEASE_NOTES = {
+      '2.0.0': 'Sample ID goes from int32 (overflowing) to int64.',
       '1.0.0': 'Initial release.',
   }
 

tensorflow_datasets/datasets/duke_ultrasound/duke_ultrasound_dataset_builder.py

@@ -16,8 +16,8 @@
 """DAS beamformed phantom images and paired clinical post-processed images."""
 
 import csv
+from datetime import datetime
 import os
-
 from etils import epath
 import numpy as np
 from tensorflow_datasets.core.utils import bool_utils
@@ -44,6 +44,7 @@ class Builder(tfds.core.GeneratorBasedBuilder):
 
   VERSION = tfds.core.Version('1.0.1')
   RELEASE_NOTES = {
+      '2.0.0': r'Fix timestamp_id from %Y%m%d%H%M%S to posix timestamp.',
       '1.0.1': 'Fixes parsing of boolean field `harmonic`.',
       '1.0.0': 'Initial release.',
   }
@@ -124,6 +125,9 @@ def _generate_examples(self, datapath, csvpath):
         iq = iq / iq.max()
         iq = 20 * np.log10(iq)
 
+        timestamp_id = datetime.strptime(row['timestamp_id'], '%Y%m%d%H%M%S')
+        timestamp = int(timestamp_id.timestamp())
+
         yield row['filename'], {
             'das': {
                 'dB': iq.astype(np.float32),
@@ -143,6 +147,6 @@ def _generate_examples(self, datapath, csvpath):
             'probe': row['probe'],
             'scanner': row['scanner'],
             'target': row['target'],
-            'timestamp_id': row['timestamp_id'],
+            'timestamp_id': timestamp,
             'harmonic': bool_utils.parse_bool(row['harm']),
         }