Files

84 lines
2.7 KiB
Python

from pathlib import Path
from uuid import uuid4
import cv2
import numpy as np
from django.conf import settings
from .algorithms import ProcessingError, ensure_uint8, image_to_data_url
def session_dir():
path = Path(settings.MEDIA_ROOT) / "sessions"
path.mkdir(parents=True, exist_ok=True)
return path
def save_image_array(image, prefix="image"):
filename = f"sessions/{prefix}-{uuid4().hex}.png"
path = Path(settings.MEDIA_ROOT) / filename
path.parent.mkdir(parents=True, exist_ok=True)
array = ensure_uint8(image)
if array.ndim == 2:
encoded_source = array
else:
encoded_source = cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
ok, encoded = cv2.imencode(".png", encoded_source)
if not ok:
raise ProcessingError("Unable to encode image for temporary storage.")
path.write_bytes(encoded.tobytes())
return filename
def save_fft_array(spectrum, prefix="fft"):
filename = f"sessions/{prefix}-{uuid4().hex}.npz"
path = Path(settings.MEDIA_ROOT) / filename
path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(path, spectrum=spectrum)
return filename
def load_image_array(relative_path):
path = Path(settings.MEDIA_ROOT) / relative_path
if not path.exists():
raise ProcessingError("Temporary image file is missing or unreadable.")
raw = np.frombuffer(path.read_bytes(), dtype=np.uint8)
image = cv2.imdecode(raw, cv2.IMREAD_UNCHANGED)
if image is None:
raise ProcessingError("Temporary image file is missing or unreadable.")
if image.ndim == 2:
return image.astype(np.uint8)
if image.shape[2] == 4:
return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB).astype(np.uint8)
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.uint8)
def load_fft_array(relative_path):
path = Path(settings.MEDIA_ROOT) / relative_path
if not path.exists():
raise ProcessingError("FFT data file is missing or unreadable.")
try:
with np.load(path) as payload:
return payload["spectrum"]
except (OSError, KeyError, ValueError) as exc:
raise ProcessingError("FFT data file is missing or unreadable.") from exc
def delete_relative_file(relative_path):
if not relative_path:
return
path = (Path(settings.MEDIA_ROOT) / relative_path).resolve()
media_root = Path(settings.MEDIA_ROOT).resolve()
if media_root not in path.parents and path != media_root:
return
if path.exists():
path.unlink()
def payload_for_image(image, relative_path):
return {
"image_path": relative_path,
"image_url": f"{settings.MEDIA_URL}{relative_path}",
"image_data": image_to_data_url(image),
}