file_storage.py

from pathlib import Path
import numpy as np
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")