Repository: Lightning-AI/pytorch-lightning
Pull request: WandbLogger's log_image can use step argument
## What does this PR do?

Makes WandbLogger's log_image compatible with the step argument.

Fixes #11703

### Does your PR introduce any breaking changes? If yes, please list them.

None

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [ ] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [ ] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)


## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
Using the step argument raises an error for WandbLogger's log_image
## 🐛 Bug

When using the WandbLogger log_image method, adding the step argument raises an error:
```
  File "/home/nathan/work/these/engine/venv/lib/python3.8/site-packages/pytorch_lightning/loggers/wandb.py", line 417, in log_image
    if len(v) != n:
TypeError: object of type 'int' has no len()
```

### To Reproduce
```

from pytorch_lightning.loggers import WandbLogger
import torch
logger=WandbLogger()
fake_img = torch.randn((50,50))
logger.log_image(key='fake', images=[fake_img], step=0)
```

### Expected behavior

The image is logged with appropriate step in Wandb.

### Environment

- PyTorch Lightning Version (e.g., 1.5.0): stable


### Additional context

Can be easily fixed by swapping a few lines in the method.


cc @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire @manangoel99

Hints from issue discussion:
hey @NathanGodey 

good catch! mind send a PR?

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py
index ee1b1b9e948be..bdb5621791936 100644
--- a/pytorch_lightning/loggers/wandb.py
+++ b/pytorch_lightning/loggers/wandb.py
@@ -413,11 +413,11 @@ def log_image(self, key: str, images: List[Any], **kwargs: str) -> None:
         """
         if not isinstance(images, list):
             raise TypeError(f'Expected a list as "images", found {type(images)}')
+        step = kwargs.pop("step", None)
         n = len(images)
         for k, v in kwargs.items():
             if len(v) != n:
                 raise ValueError(f"Expected {n} items but only found {len(v)} for {k}")
-        step = kwargs.pop("step", None)
         kwarg_list = [{k: kwargs[k][i] for k in kwargs.keys()} for i in range(n)]
         metrics = {key: [wandb.Image(img, **kwarg) for img, kwarg in zip(images, kwarg_list)]}
         self.log_metrics(metrics, step)
diff --git a/tests/loggers/test_wandb.py b/tests/loggers/test_wandb.py
index 705fe8f4b3a82..c447e8d8e92ea 100644
--- a/tests/loggers/test_wandb.py
+++ b/tests/loggers/test_wandb.py
@@ -251,6 +251,12 @@ def test_wandb_log_media(wandb, tmpdir):
     wandb.Image.assert_called_with("2.jpg")
     wandb.init().log.assert_called_once_with({"samples": [wandb.Image(), wandb.Image()]})
 
+    # test log_image with step
+    wandb.init().log.reset_mock()
+    logger.log_image(key="samples", images=["1.jpg", "2.jpg"], step=5)
+    wandb.Image.assert_called_with("2.jpg")
+    wandb.init().log.assert_called_once_with({"samples": [wandb.Image(), wandb.Image()], "trainer/global_step": 5})
+
     # test log_image with captions
     wandb.init().log.reset_mock()
     wandb.Image.reset_mock()


Code review comments:
diff_hunk: @@ -413,11 +413,11 @@ def log_image(self, key: str, images: List[Any], **kwargs: str) -> None:
         """
         if not isinstance(images, list):
             raise TypeError(f'Expected a list as "images", found {type(images)}')
+        step = kwargs.pop("step", None)
path: pytorch_lightning/loggers/wandb.py
text: @user1:
or maybe we could add `step` as a parameter (default None) and keep `kwargs` completely separate?

@user1:
also, mind adding an entry in the changelog?

@author:
I thought it did not introduce any actual change to the expected behaviour, does it make sense to add it anyway ? Like "Make step argument work properly in WandbLogger's log_image" ?

@user1:
it's a bug fix, so yes, under the bugfix section.