app/storage.py

# app/storage.py
import json
from pathlib import Path
from typing import Optional, Dict, Any

from app.configs import IMAGES_DIR, SEGMENTS_DIR, RESULTS_DIR


# =========================
# Image Storage
# =========================
def save_image(image_id: str, file_bytes: bytes, ext: str = ".jpg") -> Path:
    """Save uploaded image to storage"""
    image_path = IMAGES_DIR / f"{image_id}{ext}"
    with open(image_path, "wb") as f:
        f.write(file_bytes)
    return image_path


def get_image_path(image_id: str) -> Optional[Path]:
    """Get path of stored image by ID"""
    for ext in [".jpg", ".jpeg", ".png"]:
        path = IMAGES_DIR / f"{image_id}{ext}"
        if path.exists():
            return path
    return None


def get_image_bytes(image_id: str) -> Optional[bytes]:
    """Get image bytes by ID"""
    path = get_image_path(image_id)
    if path is None:
        return None
    with open(path, "rb") as f:
        return f.read()


# =========================
# Classification Results
# =========================
def save_classification_result(image_id: str, result: Dict[str, Any]) -> Path:
    """Save classification result as JSON"""
    result_path = RESULTS_DIR / f"{image_id}_classify.json"
    with open(result_path, "w") as f:
        json.dump(result, f, indent=2)
    return result_path


def get_classification_result(image_id: str) -> Optional[Dict[str, Any]]:
    """Get classification result by image ID"""
    result_path = RESULTS_DIR / f"{image_id}_classify.json"
    if not result_path.exists():
        return None
    with open(result_path, "r") as f:
        return json.load(f)


# =========================
# Segmentation Results
# =========================
def save_segmentation_result(image_id: str, seg_result: Dict[str, Any]) -> Optional[str]:
    """Save segmentation result as JSON, returns path"""
    result_path = RESULTS_DIR / f"{image_id}_segment.json"
    with open(result_path, "w") as f:
        json.dump(seg_result, f, indent=2)
    return str(result_path)


def get_segmentation_result(image_id: str) -> Optional[Dict[str, Any]]:
    """Get segmentation result by image ID"""
    result_path = RESULTS_DIR / f"{image_id}_segment.json"
    if not result_path.exists():
        return None
    with open(result_path, "r") as f:
        return json.load(f)