Spaces:
Runtime error
Runtime error
File size: 4,028 Bytes
edbf640 913537b edbf640 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """Storage abstraction layer for garment images and data.
Supports two backends controlled by STORAGE_BACKEND env var:
- "local" (default): filesystem at data/garments/
- "s3": HuggingFace Spaces S3 bucket for persistent storage
"""
import io
import logging
import os
from pathlib import Path
logger = logging.getLogger(__name__)
STORAGE_BACKEND = os.environ.get("STORAGE_BACKEND", "local")
LOCAL_GARMENTS_DIR = Path(__file__).resolve().parent.parent / "data" / "garments"
S3_BUCKET_NAME = os.environ.get("S3_BUCKET_NAME", "")
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", "")
S3_PREFIX = os.environ.get("S3_PREFIX", "garments/")
_s3_client = None
def _get_s3_client():
global _s3_client
if _s3_client is None:
import boto3
kwargs = {}
if S3_ENDPOINT_URL:
kwargs["endpoint_url"] = S3_ENDPOINT_URL
_s3_client = boto3.client("s3", **kwargs)
return _s3_client
def _ensure_local_dir():
LOCAL_GARMENTS_DIR.mkdir(parents=True, exist_ok=True)
def save_image(garment_id: str, image_bytes: bytes, extension: str = "jpg") -> str:
"""Save a garment image and return its storage reference.
The reference is backend-agnostic (just the filename). Use
get_image_url() to resolve it to a serveable path/URL.
"""
filename = f"{garment_id}.{extension}"
if STORAGE_BACKEND == "s3":
_save_to_s3(filename, image_bytes)
else:
_save_to_local(filename, image_bytes)
logger.info("Saved image: %s (backend: %s)", filename, STORAGE_BACKEND)
return filename
def _save_to_local(filename: str, image_bytes: bytes):
_ensure_local_dir()
path = LOCAL_GARMENTS_DIR / filename
path.write_bytes(image_bytes)
def _save_to_s3(filename: str, image_bytes: bytes):
client = _get_s3_client()
key = f"{S3_PREFIX}{filename}"
client.put_object(
Bucket=S3_BUCKET_NAME,
Key=key,
Body=image_bytes,
ContentType="image/jpeg",
)
def get_image_path(garment_id: str, extension: str = "jpg") -> str | None:
"""Return the local filesystem path for a garment image.
For S3 backend, downloads to a local cache first.
Returns None if the image doesn't exist.
"""
filename = f"{garment_id}.{extension}"
if STORAGE_BACKEND == "s3":
return _download_from_s3(filename)
path = LOCAL_GARMENTS_DIR / filename
if path.exists():
return str(path)
return None
def _download_from_s3(filename: str) -> str | None:
_ensure_local_dir()
local_path = LOCAL_GARMENTS_DIR / filename
if local_path.exists():
return str(local_path)
try:
client = _get_s3_client()
key = f"{S3_PREFIX}{filename}"
response = client.get_object(Bucket=S3_BUCKET_NAME, Key=key)
local_path.write_bytes(response["Body"].read())
return str(local_path)
except Exception as e:
logger.warning("Failed to download %s from S3: %s", filename, e)
return None
def delete_image(garment_id: str, extension: str = "jpg"):
"""Remove a garment image from storage."""
filename = f"{garment_id}.{extension}"
if STORAGE_BACKEND == "s3":
try:
client = _get_s3_client()
key = f"{S3_PREFIX}{filename}"
client.delete_object(Bucket=S3_BUCKET_NAME, Key=key)
except Exception as e:
logger.warning("Failed to delete %s from S3: %s", filename, e)
local_path = LOCAL_GARMENTS_DIR / filename
if local_path.exists():
local_path.unlink()
def image_exists(garment_id: str, extension: str = "jpg") -> bool:
"""Check if a garment image exists in storage."""
filename = f"{garment_id}.{extension}"
if STORAGE_BACKEND == "s3":
try:
client = _get_s3_client()
key = f"{S3_PREFIX}{filename}"
client.head_object(Bucket=S3_BUCKET_NAME, Key=key)
return True
except Exception:
return False
return (LOCAL_GARMENTS_DIR / filename).exists()
|