emotion-detection-api / app /image_storage.py
HimAJ's picture
upload 32 files for the ml
1e4fc28 verified
"""
Image storage utilities for saving and serving uploaded images.
"""
import os
import uuid
import shutil
from pathlib import Path
from typing import Optional, Tuple
from werkzeug.utils import secure_filename
def ensure_images_dir(images_dir: str) -> str:
"""Ensure images directory exists and return its path."""
os.makedirs(images_dir, exist_ok=True)
return images_dir
def generate_unique_filename(original_filename: str) -> str:
"""
Generate a unique filename to avoid collisions.
Format: {uuid}_{secure_original_name} or just {uuid}.jpg if original is invalid
"""
# Get secure base name
base_name = secure_filename(original_filename)
if not base_name:
base_name = "upload.jpg"
# Add UUID prefix for uniqueness (use full UUID to ensure uniqueness)
name, ext = os.path.splitext(base_name)
if not ext or ext.lower() not in ('.jpg', '.jpeg', '.png'):
ext = '.jpg'
unique_id = str(uuid.uuid4()) # Full UUID for better uniqueness
return f"{unique_id}_{name}{ext}"
def save_image(source_path: str, images_dir: str, original_filename: str) -> Optional[str]:
"""
Save an image from source_path to images_dir with a unique filename.
Args:
source_path: Path to source image file
images_dir: Directory to save images to
original_filename: Original filename for reference
Returns:
Stored filename (relative to images_dir) or None on failure
"""
try:
ensure_images_dir(images_dir)
# Generate unique filename
stored_filename = generate_unique_filename(original_filename)
dest_path = os.path.join(images_dir, stored_filename)
# Copy file
shutil.copy2(source_path, dest_path)
return stored_filename
except Exception as e:
# Log error but don't fail the request
import logging
logging.getLogger(__name__).exception(f"Failed to save image: {e}")
return None
def get_image_path(images_dir: str, filename: str) -> Optional[str]:
"""
Get full path to an image file if it exists.
Args:
images_dir: Base images directory
filename: Image filename
Returns:
Full path to image or None if not found
"""
if not filename:
return None
# Security: ensure filename doesn't contain path traversal
# Extract just the basename to prevent directory traversal
base_filename = os.path.basename(filename)
safe_filename = secure_filename(base_filename)
if not safe_filename:
return None
# Use safe_filename for the path (secure_filename may have sanitized it)
# But also try the original if it's already safe
image_path = os.path.join(images_dir, safe_filename)
if os.path.exists(image_path) and os.path.isfile(image_path):
return image_path
# Also try the original filename if it's different and seems safe
if safe_filename != base_filename:
# Check if original is safe (no path separators, no parent dir references)
if base_filename == filename and '/' not in base_filename and '\\' not in base_filename and '..' not in base_filename:
alt_path = os.path.join(images_dir, base_filename)
if os.path.exists(alt_path) and os.path.isfile(alt_path):
return alt_path
return None
def delete_image(images_dir: str, filename: str) -> bool:
"""
Delete an image file.
Args:
images_dir: Base images directory
filename: Image filename to delete
Returns:
True if deleted, False otherwise
"""
try:
image_path = get_image_path(images_dir, filename)
if image_path and os.path.exists(image_path):
os.remove(image_path)
return True
return False
except Exception:
return False