File size: 2,783 Bytes
aac350d | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """
Filesystem artifact store.
Responsible for saving uploaded source images, generated annotated
images, and any other binary artifacts. Returns relative paths so the
API can construct download URLs without coupling to the filesystem layout.
"""
from __future__ import annotations
import time
from pathlib import Path
from typing import Optional
import cv2
import numpy as np
from loguru import logger
class ArtifactStore:
"""Saves and reads image artifacts under a root directory."""
def __init__(self, root: Path) -> None:
self._root = Path(root)
self._uploads = self._root / "uploads"
self._generated = self._root / "generated"
self._uploads.mkdir(parents=True, exist_ok=True)
self._generated.mkdir(parents=True, exist_ok=True)
# ------------------------------------------------------------------ #
# Uploads (source images from users)
# ------------------------------------------------------------------ #
def save_upload(self, img: np.ndarray, prefix: str = "upload") -> Path:
"""Save an uploaded image. Returns absolute path."""
path = self._uploads / f"{prefix}_{int(time.time() * 1000)}.jpg"
cv2.imwrite(str(path), img)
return path
def save_upload_bytes(self, data: bytes, prefix: str = "upload",
suffix: str = ".jpg") -> Path:
path = self._uploads / f"{prefix}_{int(time.time() * 1000)}{suffix}"
path.write_bytes(data)
return path
# ------------------------------------------------------------------ #
# Generated artifacts (annotated images, crops, montages)
# ------------------------------------------------------------------ #
def save_generated(self, img: np.ndarray, prefix: str = "gen") -> Path:
path = self._generated / f"{prefix}_{int(time.time() * 1000)}.jpg"
cv2.imwrite(str(path), img)
return path
# ------------------------------------------------------------------ #
# Lookup
# ------------------------------------------------------------------ #
def resolve(self, relative_path: str) -> Optional[Path]:
"""Resolve a relative path to absolute, with root-jail safety."""
target = (self._root / relative_path).resolve()
try:
target.relative_to(self._root.resolve())
except ValueError:
logger.warning(f"Path traversal attempt blocked: {relative_path}")
return None
if not target.exists():
return None
return target
@property
def root(self) -> Path:
return self._root
@property
def uploads_dir(self) -> Path:
return self._uploads
@property
def generated_dir(self) -> Path:
return self._generated
|