Spaces:
Running on Zero
Running on Zero
| from pathlib import Path | |
| import shutil | |
| from typing import Any, Optional | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from pytorch_lightning.loggers.logger import Logger | |
| from pytorch_lightning.utilities import rank_zero_only | |
| class LocalLogger(Logger): | |
| def __init__(self, log_path: Path | str = Path("outputs/local")) -> None: | |
| super().__init__() | |
| self.experiment = None | |
| self.log_path = Path(log_path) | |
| shutil.rmtree(self.log_path, ignore_errors=True) | |
| def name(self): | |
| return "LocalLogger" | |
| def version(self): | |
| return 0 | |
| def log_hyperparams(self, params): | |
| pass | |
| def log_metrics(self, metrics, step): | |
| pass | |
| def log_image( | |
| self, | |
| key: str, | |
| images: list[Any], | |
| step: Optional[int] = None, | |
| **kwargs, | |
| ): | |
| # The function signature is the same as the wandb logger's, but the step is | |
| # actually required. | |
| assert step is not None | |
| for index, image in enumerate(images): | |
| path = self.log_path / f"{key}/{index:0>2}_{step:0>6}.png" | |
| path.parent.mkdir(exist_ok=True, parents=True) | |
| if isinstance(image, torch.Tensor): | |
| Image.fromarray(image.permute(1, 2, 0).numpy().astype(np.uint8)).save(path) | |
| else: | |
| Image.fromarray(image).save(path) | |