test-image / file_storage.py
abidlabs's picture
abidlabs HF Staff
Upload folder using huggingface_hub
f5a2c54 verified
from pathlib import Path
from PIL import Image as PILImage
try: # absolute imports when installed
from trackio.utils import TRACKIO_DIR
except ImportError: # relative imports for local execution on Spaces
from utils import TRACKIO_DIR
class FileStorage:
@staticmethod
def get_project_media_path(
project: str,
run: str | None = None,
step: int | None = None,
filename: str | None = None,
) -> Path:
if filename is not None and step is None:
raise ValueError("filename requires step")
if step is not None and run is None:
raise ValueError("step requires run")
path = TRACKIO_DIR / "media" / project
if run:
path /= run
if step is not None:
path /= str(step)
if filename:
path /= filename
return path
@staticmethod
def init_project_media_path(
project: str, run: str | None = None, step: int | None = None
) -> Path:
path = FileStorage.get_project_media_path(project, run, step)
path.mkdir(parents=True, exist_ok=True)
return path
@staticmethod
def save_image(
image: PILImage.Image,
project: str,
run: str,
step: int,
filename: str,
format: str = "PNG",
) -> Path:
path = FileStorage.init_project_media_path(project, run, step) / filename
image.save(path, format=format)
return path
@staticmethod
def get_image(project: str, run: str, step: int, filename: str) -> PILImage.Image:
path = FileStorage.get_project_media_path(project, run, step, filename)
if not path.exists():
raise FileNotFoundError(f"Image file not found: {path}")
return PILImage.open(path).convert("RGBA")