feat: Support processing RAW format images #17

This commit is contained in:
shiyu
2025-08-25 17:37:50 +08:00
parent a0361db82a
commit 3d0bdd6557
5 changed files with 92 additions and 13 deletions

View File

@@ -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:

View File

@@ -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

View File

@@ -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

View File

@@ -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)

View File

@@ -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,