Spaces:
Sleeping
Sleeping
File size: 3,337 Bytes
fc895f4 | 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 | """
loader.py
---------
Handles loading floor plan images from disk.
Supports: PNG, JPG, JPEG, BMP, TIFF, PDF (first page via PIL).
Normalizes to a standard resolution while preserving aspect ratio.
"""
import cv2
import numpy as np
from PIL import Image
from pathlib import Path
SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pdf"}
def load_image(image_path: str, target_size: int = 1024) -> np.ndarray:
"""
Load a floor plan image and normalize it to a standard size.
Args:
image_path: Path to the input image file.
target_size: The longer dimension will be resized to this value.
Aspect ratio is preserved.
Returns:
Grayscale numpy array of shape (H, W), dtype uint8.
Raises:
FileNotFoundError: If the file does not exist.
ValueError: If the file format is not supported.
"""
path = Path(image_path)
if not path.exists():
raise FileNotFoundError(f"Image not found: {image_path}")
if path.suffix.lower() not in SUPPORTED_FORMATS:
raise ValueError(
f"Unsupported format '{path.suffix}'. "
f"Supported: {SUPPORTED_FORMATS}"
)
# PDF: extract first page as image
if path.suffix.lower() == ".pdf":
img = _load_pdf_page(path)
else:
img = _load_raster(path)
# Convert to grayscale if needed
if len(img.shape) == 3:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Resize to target_size along the longer dimension
img = _resize_keep_aspect(img, target_size)
return img
def _load_raster(path: Path) -> np.ndarray:
"""Load PNG/JPG/BMP/TIFF using OpenCV."""
img = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
if img is None:
raise IOError(f"OpenCV could not read image: {path}")
return img
def _load_pdf_page(path: Path, page: int = 0, dpi: int = 200) -> np.ndarray:
"""
Load the first page of a PDF as a numpy array using PIL.
Requires Pillow with PDF support.
"""
try:
pil_img = Image.open(str(path))
pil_img.load()
# For multi-page PDFs, seek to desired page
if hasattr(pil_img, "n_frames") and pil_img.n_frames > page:
pil_img.seek(page)
pil_img = pil_img.convert("RGB")
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
except Exception as e:
raise IOError(f"Failed to load PDF '{path}': {e}")
def _resize_keep_aspect(img: np.ndarray, target_size: int) -> np.ndarray:
"""
Resize image so its longer dimension equals target_size.
Uses INTER_AREA for downscaling (best quality for line drawings).
"""
h, w = img.shape[:2]
if max(h, w) == target_size:
return img
scale = target_size / max(h, w)
new_w = int(w * scale)
new_h = int(h * scale)
interpolation = cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR
return cv2.resize(img, (new_w, new_h), interpolation=interpolation)
def save_image(img: np.ndarray, output_path: str) -> None:
"""Save a numpy array as an image file."""
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
cv2.imwrite(output_path, img)
print(f"Saved: {output_path}")
|