From 3d0bdd6557819f8e679dab3b68b7555d94facb13 Mon Sep 17 00:00:00 2001 From: shiyu Date: Mon, 25 Aug 2025 17:37:50 +0800 Subject: [PATCH] feat: Support processing RAW format images #17 --- api/routes/virtual_fs.py | 22 +++++++++++++-- requirements.txt | 2 ++ services/thumbnail.py | 47 +++++++++++++++++++++++++------ services/virtual_fs.py | 32 ++++++++++++++++++++- web/src/apps/ImageViewer/index.ts | 2 +- 5 files changed, 92 insertions(+), 13 deletions(-) diff --git a/api/routes/virtual_fs.py b/api/routes/virtual_fs.py index fd1c840..2efa1e5 100644 --- a/api/routes/virtual_fs.py +++ b/api/routes/virtual_fs.py @@ -16,11 +16,9 @@ from services.virtual_fs import ( generate_temp_link_token, verify_temp_link_token, ) -from services.thumbnail import is_image_filename, get_or_create_thumb +from services.thumbnail import is_image_filename, get_or_create_thumb, is_raw_filename from schemas import MkdirRequest, MoveRequest from api.response import success -from services.ai import get_text_embedding -from services.vector_db import VectorDBService router = APIRouter(prefix='/api/fs', tags=["virtual-fs"]) @@ -33,6 +31,24 @@ async def get_file( ): full_path = '/' + full_path if not full_path.startswith('/') else full_path + if is_raw_filename(full_path): + import rawpy + from PIL import Image + import io + try: + raw_data = await read_file(full_path) + with rawpy.imread(io.BytesIO(raw_data)) as raw: + rgb = raw.postprocess(use_camera_wb=True, output_bps=8) + im = Image.fromarray(rgb) + buf = io.BytesIO() + im.save(buf, 'JPEG', quality=90) + content = buf.getvalue() + return Response(content=content, media_type='image/jpeg') + except FileNotFoundError: + raise HTTPException(404, detail="File not found") + except Exception as e: + raise HTTPException(500, detail=f"RAW file processing failed: {e}") + try: content = await read_file(full_path) except FileNotFoundError: diff --git a/requirements.txt b/requirements.txt index 5872af3..d11bf02 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,6 +15,7 @@ httpcore==1.0.9 httptools==0.6.4 httpx==0.28.1 idna==3.10 +imageio==2.37.0 iso8601==2.1.0 Jinja2==3.1.6 markdown-it-py==4.0.0 @@ -40,6 +41,7 @@ python-dotenv==1.1.1 python-multipart==0.0.20 pytz==2025.2 PyYAML==6.0.2 +rawpy==0.25.1 rich==14.1.0 rich-toolkit==0.15.0 rignore==0.6.4 diff --git a/services/thumbnail.py b/services/thumbnail.py index db8ceb1..94609fc 100644 --- a/services/thumbnail.py +++ b/services/thumbnail.py @@ -5,7 +5,8 @@ from pathlib import Path from typing import Tuple from fastapi import HTTPException -ALLOWED_EXT = {"jpg", "jpeg", "png", "webp", "gif", "bmp", "tiff"} +ALLOWED_EXT = {"jpg", "jpeg", "png", "webp", "gif", "bmp", "tiff", "arw", "cr2", "cr3", "nef", "rw2", "orf", "pef", "dng"} +RAW_EXT = {"arw", "cr2", "cr3", "nef", "rw2", "orf", "pef", "dng"} MAX_SOURCE_SIZE = 200 * 1024 * 1024 CACHE_ROOT = Path('data/.thumb_cache') @@ -17,6 +18,13 @@ def is_image_filename(name: str) -> bool: return parts[1].lower() in ALLOWED_EXT +def is_raw_filename(name: str) -> bool: + parts = name.rsplit('.', 1) + if len(parts) < 2: + return False + return parts[1].lower() in RAW_EXT + + def _cache_key(adapter_id: int, rel: str, size: int, mtime: int, w: int, h: int, fit: str) -> str: raw = f"{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode() return hashlib.sha1(raw).hexdigest() @@ -31,9 +39,29 @@ def _ensure_cache_dir(p: Path): p.parent.mkdir(parents=True, exist_ok=True) -def generate_thumb(data: bytes, w: int, h: int, fit: str) -> Tuple[bytes, str]: +def generate_thumb(data: bytes, w: int, h: int, fit: str, is_raw: bool = False) -> Tuple[bytes, str]: from PIL import Image - im = Image.open(io.BytesIO(data)) + if is_raw: + try: + import rawpy + with rawpy.imread(io.BytesIO(data)) as raw: + try: + thumb = raw.extract_thumb() + except rawpy.LibRawNoThumbnailError: + thumb = None + + if thumb is not None and thumb.format in [rawpy.ThumbFormat.JPEG, rawpy.ThumbFormat.BITMAP]: + im = Image.open(io.BytesIO(thumb.data)) + else: + rgb = raw.postprocess(use_camera_wb=False, use_auto_wb=True, output_bps=8) + im = Image.fromarray(rgb) + except Exception as e: + print(f"rawpy processing failed: {e}") + raise e + + else: + im = Image.open(io.BytesIO(data)) + if im.mode not in ("RGB", "RGBA"): im = im.convert("RGBA" if im.mode in ("P", "LA") else "RGB") if fit == 'cover': @@ -57,17 +85,20 @@ def generate_thumb(data: bytes, w: int, h: int, fit: str) -> Tuple[bytes, str]: async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: int, h: int, fit: str = 'cover'): - read_data = await adapter.read_file(root, rel) - if len(read_data) > MAX_SOURCE_SIZE: - raise HTTPException(404, detail="Image too large for thumbnail") - key = _cache_key(adapter_id, rel, len(read_data), 0, w, h, fit) + stat = await adapter.stat_file(root, rel) + if stat['size'] > MAX_SOURCE_SIZE: + raise HTTPException(400, detail="Image too large for thumbnail") + + key = _cache_key(adapter_id, rel, stat['size'], int(stat['mtime']), w, h, fit) path = _cache_path(key) if path.exists(): return path.read_bytes(), 'image/webp', key _ensure_cache_dir(path) + read_data = await adapter.read_file(root, rel) try: - thumb_bytes, mime = generate_thumb(read_data, w, h, fit) + thumb_bytes, mime = generate_thumb(read_data, w, h, fit, is_raw=is_raw_filename(rel)) except Exception as e: + print(e) raise HTTPException(500, detail=f"Thumbnail generation failed: {e}") path.write_bytes(thumb_bytes) return thumb_bytes, mime, key diff --git a/services/virtual_fs.py b/services/virtual_fs.py index 7d19a5c..331a741 100644 --- a/services/virtual_fs.py +++ b/services/virtual_fs.py @@ -11,7 +11,7 @@ import base64 from models import Mount from .adapters.registry import runtime_registry from api.response import page -from .thumbnail import is_image_filename +from .thumbnail import is_image_filename, is_raw_filename from services.processors.registry import get as get_processor from services.tasks import task_service from services.logging import LogService @@ -341,6 +341,36 @@ async def stream_file(path: str, range_header: str | None): adapter, mount, root, rel = await resolve_adapter_and_rel(path) if not rel or rel.endswith('/'): raise HTTPException(400, detail="Path is a directory") + if is_raw_filename(rel): + import rawpy + from PIL import Image + import io + try: + raw_data = await read_file(path) + try: + import rawpy + with rawpy.imread(io.BytesIO(raw_data)) as raw: + try: + thumb = raw.extract_thumb() + except rawpy.LibRawNoThumbnailError: + thumb = None + + if thumb is not None and thumb.format in [rawpy.ThumbFormat.JPEG, rawpy.ThumbFormat.BITMAP]: + im = Image.open(io.BytesIO(thumb.data)) + else: + rgb = raw.postprocess(use_camera_wb=False, use_auto_wb=True, output_bps=8) + im = Image.fromarray(rgb) + except Exception as e: + print(f"rawpy processing failed: {e}") + raise e + + buf = io.BytesIO() + im.save(buf, 'JPEG', quality=90) + content = buf.getvalue() + return Response(content=content, media_type='image/jpeg') + except Exception as e: + raise HTTPException(500, detail=f"RAW file processing failed: {e}") + stream_impl = getattr(adapter, "stream_file", None) if callable(stream_impl): return await stream_impl(root, rel, range_header) diff --git a/web/src/apps/ImageViewer/index.ts b/web/src/apps/ImageViewer/index.ts index b688b5f..3042727 100644 --- a/web/src/apps/ImageViewer/index.ts +++ b/web/src/apps/ImageViewer/index.ts @@ -7,7 +7,7 @@ export const descriptor: AppDescriptor = { supported: (entry) => { if (entry.is_dir) return false; const ext = entry.name.split('.').pop()?.toLowerCase() || ''; - return ['png', 'jpg', 'jpeg', 'gif', 'webp', 'svg', 'bmp', 'ico', 'avif'].includes(ext); + return ['png', 'jpg', 'jpeg', 'gif', 'webp', 'svg', 'bmp', 'ico', 'avif', 'arw', 'cr2', 'cr3', 'nef', 'rw2', 'orf', 'pef', 'dng'].includes(ext); }, component: ImageViewerApp, default: true,