WandbLogger's log_image can use step argument (#11716)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed the format of the configuration saved automatically by the CLI's `SaveConfigCallback` ([#11532](https://github.com/PyTorchLightning/pytorch-lightning/pull/11532))
 - Fixed an issue to avoid validation loop run on restart ([#11552](https://github.com/PyTorchLightning/pytorch-lightning/pull/11552))
 - The Rich progress bar now correctly shows the `on_epoch` logged values on train epoch end ([#11689](https://github.com/PyTorchLightning/pytorch-lightning/pull/11689))
+- Fixed an issue to make the `step` argument in `WandbLogger.log_image` work ([#11716](https://github.com/PyTorchLightning/pytorch-lightning/pull/11716))
 
 
 ## [1.5.9] - 2022-01-18

pytorch_lightning/loggers/wandb.py

@@ -405,7 +405,7 @@ def log_text(
         self.log_table(key, columns, data, dataframe, step)
 
     @rank_zero_only
-    def log_image(self, key: str, images: List[Any], **kwargs: str) -> None:
+    def log_image(self, key: str, images: List[Any], step: Optional[int] = None, **kwargs: str) -> None:
         """Log images (tensors, numpy arrays, PIL Images or file paths).
 
         Optional kwargs are lists passed to each image (ex: caption, masks, boxes).
@@ -416,7 +416,6 @@ def log_image(self, key: str, images: List[Any], **kwargs: str) -> None:
         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)

tests/loggers/test_wandb.py

@@ -252,6 +252,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()