Spaces:
Build error
Build error
File size: 1,364 Bytes
8c3e275 | 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 | from __future__ import annotations
from pathlib import Path
import cv2
import numpy as np
from PIL import Image
IMAGE_EXTENSIONS = {
".jpg",
".jpeg",
".png",
".bmp",
".tiff",
".tif",
".webp",
".ppm",
".pgm",
".pbm",
}
def load_image(path: str | Path) -> np.ndarray:
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"Image not found: {path}")
ext = path.suffix.lower()
if ext == ".pdf":
return _load_pdf(path)
img = cv2.imread(str(path))
if img is not None:
return img
pil_img = Image.open(path)
arr = np.array(pil_img)
if arr.ndim == 2:
return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
if arr.shape[-1] == 4:
return cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR)
return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
def load_pdf_pages(path: Path, scale: float = 2.0) -> list[np.ndarray]:
import pypdfium2 as pdfium
pdf = pdfium.PdfDocument(str(path))
pages = []
for page in pdf:
bitmap = page.render(scale=scale)
pil = bitmap.to_pil()
arr = np.array(pil)
pages.append(cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
return pages
def _load_pdf(path: Path) -> np.ndarray:
pages = load_pdf_pages(path)
if pages:
return pages[0]
raise ValueError(f"PDF has no pages: {path}")
|