Spaces:
Runtime error
Runtime error
File size: 1,823 Bytes
b77494e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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")
|