[Feature Request]: Channels-first support for WandBImageClassificationCallback

[soumik12345]

Request from Ian Stenbit to make the WandBImageClassificationCallback work with both channels-first and channels-last image formats.

[mratanusarkar]

Hi @soumik12345, can I work on this issue?

I figured out that the class WandBImageClassificationCallback exists in the wandb_addons\keras\image_classification.py,
and to work with both channels-first and channels-last image formats, we need to ensure that the image tensor is correctly formatted before logging it to WandB.

[soumik12345]

@mratanusarkar Thanks for your interest. Please feel free to raise a PR.

[soumik12345]

@mratanusarkar Here are 2 examples to aid your testing:

• Use this example to test for the channels-first pipeline.
• Use this example to test for the channels-last pipeline.

[mratanusarkar]

@soumik12345 thanks!

[mratanusarkar]

@soumik12345 regarding the channels-first notebook you shared,

• The import from keras_core.utils import TorchModuleWrapper throws the following error:

  ImportError: cannot import name 'TorchModuleWrapper' from 'keras_core.utils' (/usr/local/lib/python3.10/dist-packages/keras_core/utils/__init__.py)

  Which was resolved with from keras_core.layers import TorchModuleWrapper as per your suggestion.
• The notebook doesn't use WandBImageClassificationCallback hence I am unable to test out the channels-first usecase.

Could you please share me a notebook with channels_first image data using the WandBImageClassificationCallback?
It would be of great help. thanks!

[mratanusarkar]

• Use this example to test for the channels-first pipeline.

@soumik12345 I am unable to test the channels-first pipeline.

I feel WandBImageClassificationCallback callback is designed to work with TensorFlow datasets (tf.data.Dataset) or NumPy arrays (np.array). And not with the PyTorch's DataLoader.

more details

Here are some inputs I was using to debug:

1. For channels-last notebook:

   Use this example to test for the channels-last pipeline.

   • type(test_dataset) --> tensorflow.python.data.ops.prefetch_op._PrefetchDataset
   • type(config.classes) --> list

2. For channels-first notebook:

   Use this example to test for the channels-first pipeline.

   • type(image_datasets) --> dict
   • type(image_datasets["val"]) --> torchvision.datasets.folder.ImageFolder
   • type(class_names) --> list

Hence, I believe there was no WandBImageClassificationCallback in the channels-first notebook, and hence I made the above conclusion from a little bit of debugging.

I believe this could be a good feature request.
As per the current issue, please help me with how I should proceed and close the PR.