AA.composed.ImgOptical fails with Tensorflow dataset

[wjpearson]

AA.composed.ImgOptical requires a dataset to sample from for AA.image_domain.optical.CroppedTemplateOverlap. This indexes the dataset to pull a random galaxy. Prefetched Tensorflow datasets cannot be indexed so requesting a random image fails.

I would suggest having the dataset an optional requirement in AA.composed.ImgOptical and either skipping the call to AA.image_domain.optical.CroppedTemplateOverlap or calling it with a template if dataset == None

Full traceback:

<ipython-input-7-cf2b44823a21> in <module>
     30 train_iter = iter(train)
     31 
---> 32 img = next(train_iter)['image']
     33 print(img.shape)

~/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py in __next__(self)
    759   def __next__(self):
    760     try:
--> 761       return self._next_internal()
    762     except errors.OutOfRangeError:
    763       raise StopIteration

~/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py in _next_internal(self)
    745           self._iterator_resource,
    746           output_types=self._flat_output_types,
--> 747           output_shapes=self._flat_output_shapes)
    748 
    749       try:

~/.local/lib/python3.6/site-packages/tensorflow/python/ops/gen_dataset_ops.py in iterator_get_next(iterator, output_types, output_shapes, name)
   2726       return _result
   2727     except _core._NotOkStatusException as e:
-> 2728       _ops.raise_from_not_ok_status(e, name)
   2729     except _core._FallbackException:
   2730       pass

~/.local/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6939   message = e.message + (" name: " + name if name is not None else "")
   6940   # pylint: disable=protected-access
-> 6941   six.raise_from(core._status_to_exception(e.code, message), None)
   6942   # pylint: enable=protected-access
   6943 

~/.local/lib/python3.6/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: TypeError: 'PrefetchDataset' object does not support indexing
Traceback (most recent call last):

  File "/home/pearsonw/.local/lib/python3.6/site-packages/tensorflow/python/ops/script_ops.py", line 249, in __call__
    ret = func(*args)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 645, in wrapper
    return func(*args, **kwargs)

  File "<ipython-input-7-cf2b44823a21>", line 13, in aug_fn
    aug_data = transform(**data)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/composition.py", line 205, in __call__
    data = t(force_apply=force_apply, **data)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 95, in __call__
    return self.apply_with_params(params, **kwargs)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 110, in apply_with_params
    res[key] = target_function(arg, **dict(params, **target_dependencies))

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/augmentations/transforms.py", line 2219, in apply
    return fn(img, **params)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/astroaugmentations/composed.py", line 183, in __call__
    return self.augmentation(image=image)['image']

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/composition.py", line 205, in __call__
    data = t(force_apply=force_apply, **data)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 95, in __call__
    return self.apply_with_params(params, **kwargs)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 110, in apply_with_params
    res[key] = target_function(arg, **dict(params, **target_dependencies))

  File "/home/pearsonw/.local/lib/python3.6/site-packages/albumentations/augmentations/transforms.py", line 2219, in apply
    return fn(img, **params)

  File "/home/pearsonw/.local/lib/python3.6/site-packages/astroaugmentations/image_domain/optical.py", line 310, in __call__
    datasample, _ = self.dataset[self.rng.integers(0, len(self.dataset))]

TypeError: 'PrefetchDataset' object does not support indexing


	 [[{{node augment}}]] [Op:IteratorGetNext]

(I did promise you I would break it)

[mb010]

Good point. I didn't think of this. For now, I can disable this component of that composed augmentation if not explicitly requested.
If the data template were used as an iterable, instead of an indexed item, would that work with TF?
And thanks for breaking it! 😄

[wjpearson]

I think using it as an iterable should work in TF.