Spaces:
Sleeping
Sleeping
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: | |
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 | |
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 | |
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 | |
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") | |