[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](https://colab.research.google.com/github/soumik12345/wandb-addons/blob/main/docs/keras/examples/image_classification.ipynb) 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](https://colab.research.google.com/github/wandb/examples/blob/example/torchvision-keras/colabs/keras/keras_core/torchvision-keras.ipynb) 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.