File size: 5,194 Bytes
46cf14b d18bae1 |
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 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import io
import json
import shutil
import sqlite3
from pathlib import Path
from fastapi import APIRouter, UploadFile, File, Query, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from storage.files.file_manager import FileManager
from storage.common import validate_token
router = APIRouter(prefix="/embeddings", tags=["Embeddings Manager"])
EMBEDDINGS_ROOT = Path("/data/embeddings")
file_manager = FileManager(EMBEDDINGS_ROOT)
HF_TOKEN = os.getenv("HF_TOKEN")
@router.get("/list_embeddings", tags=["Embeddings Manager"])
def list_all_embeddings(
token: str = Query(..., description="Token required for authorization")
):
"""
List all embeddings stored under /data/embeddings.
For each video hash folder, returns:
- video: folder name (hash)
- faces: true/false depending on whether faces/embeddings.json exists
- voices: true/false depending on whether voices/embeddings.json exists
Notes:
- A video folder may contain only faces, only voices, or neither.
- Missing folders are treated as false.
"""
validate_token(token)
results = []
# If embeddings root does not exist, return empty list
if not EMBEDDINGS_ROOT.exists():
return []
for video_dir in EMBEDDINGS_ROOT.iterdir():
if not video_dir.is_dir():
continue # Skip anything that is not a folder
faces_path = video_dir / "faces" / "embeddings.json"
voices_path = video_dir / "voices" / "embeddings.json"
results.append({
"video": video_dir.name,
"faces": faces_path.exists(),
"voices": voices_path.exists()
})
return results
@router.post("/upload_embeddings", tags=["Embeddings Manager"])
async def upload_embeddings(
file: UploadFile = File(...),
embedding_type: str = Query(..., description="faces or voices"),
video_hash: str = Query(..., description="Hash of the video"),
token: str = Query(..., description="Token required for authorization")
):
"""
Upload embeddings JSON for a given video and type (faces or voices).
Behavior:
- Validate the token.
- Validate embedding_type.
- Ensure directory structure: /data/embeddings/<video_hash>/<embedding_type>/
- Delete any existing .json file inside that folder.
- Save the uploaded embeddings as embeddings.json.
"""
validate_token(token)
# Validación del tipo
if embedding_type not in ("faces", "voices"):
raise HTTPException(status_code=400, detail="embedding_type must be 'faces' or 'voices'")
# Rutas objetivo
video_path = EMBEDDINGS_ROOT / video_hash
type_path = video_path / embedding_type
# Crear carpetas si no existen
type_path.mkdir(parents=True, exist_ok=True)
# Eliminar JSONs previos
for existing in type_path.glob("*.json"):
try:
existing.unlink()
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Failed to delete old embeddings: {exc}")
# Guardar como embeddings.json
final_path = type_path / "embeddings.json"
try:
file_bytes = await file.read()
with open(final_path, "wb") as f:
f.write(file_bytes)
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Failed to save embeddings: {exc}")
return JSONResponse(
status_code=200,
content={
"status": "ok",
"saved_to": str(final_path)
}
)
def get_embeddings_json(video_hash: str, embedding_type: str):
"""
Returns the parsed embeddings.json for a given video and type.
Behavior:
- Validate embedding_type.
- Build the file path: /data/embeddings/<video_hash>/<embedding_type>/embeddings.json
- Raise HTTPException if missing.
- Load and return parsed JSON.
"""
if embedding_type not in ("faces", "voices"):
raise HTTPException(status_code=400, detail="embedding_type must be 'faces' or 'voices'")
target_file = EMBEDDINGS_ROOT / video_hash / embedding_type / "embeddings.json"
if not target_file.exists():
raise HTTPException(
status_code=404,
detail=f"embeddings.json not found for video={video_hash}, type={embedding_type}"
)
try:
with open(target_file, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Failed to read embeddings: {exc}")
return data
@router.get("/get_embedding", tags=["Embeddings Manager"])
def get_embeddings(
video_hash: str = Query(..., description="Hash of the video"),
embedding_type: str = Query(..., description="faces or voices"),
token: str = Query(..., description="Token required for authorization")
):
"""
Endpoint to retrieve embeddings.json for a given video hash and type.
"""
validate_token(token)
data = get_embeddings_json(video_hash, embedding_type)
return data |