okooooooo/Pixhelf
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add semantic image and text search

Browse Source
This commit is contained in:
adminroot committed 2026-08-29 05:30:16 +00:00
1 parent 6a9164e5ea
commit 77a8d7ed66
17 files changed
+6152 -85
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+1072 -3
View File
File diff suppressed because it is too large. Load diff
Cargo.toml
+4
View File
@@ -8,6 +8,9 @@ build = "build.rs"
anyhow = "1.0.104"
axum = { version = "0.8.9", features = ["json", "query", "tokio"] }
blake3 = "1.8.7"
candle-core = { version = "0.9.2", default-features = false }
candle-nn = { version = "0.9.2", default-features = false }
candle-transformers = { version = "0.9.2", default-features = false }
image = { version = "0.25.10", default-features = false, features = ["jpeg", "png", "webp"] }
natord = "1.0.9"
serde = { version = "1.0.229", features = ["derive"] }
@@ -17,6 +20,7 @@ tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.11", features = ["compression-br", "compression-gzip", "fs", "set-header", "trace"] }
tracing = "0.1.44"
tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }
tokenizers = { version = "0.22.2", default-features = false, features = ["fancy-regex"] }
walkdir = "2.5.0"
webp = "0.3.1"
README.md
+86 -1
View File
@@ -1,6 +1,6 @@
# Pixhelf
轻量的本地照片画廊，支持 JPEG、PNG、WebP、搜索和排序。
轻量的本地照片画廊，支持 JPEG、PNG、WebP、自然语言文字搜图、随机探索和本地以图搜图。相似图片会直接显示在图片查看器下方的分页瀑布流中。
## Docker Compose
@@ -31,3 +31,88 @@ mkdir -p pic && docker run -d --name pixhelf --restart unless-stopped -p 3002:30
```
照片放进 `./pic`，访问 <http://localhost:3002>。
## 自然语言文字搜图
可选的 Chinese-CLIP ViT-B/16 模型支持直接输入中文描述，例如“海边日落”“草地上的狗”或“夜晚城市街道”。图片与文字都在本机编码，Pixhelf 不会在运行时联网或自动下载模型。索引尚未完成时搜索框仍按文件名搜索；完成后自动切换到语义搜索，并保留文件名精确匹配优先级。
下载并校验固定版本的官方模型与词表（模型约 753 MB）：
```bash
mkdir -p models/chinese-clip
curl --proto '=https' --tlsv1.2 --fail --location \
https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/resolve/f4a64596bbcf9a2a94591b74b9dc39b2e4e77e3e/model.safetensors \
--output models/chinese-clip/model.safetensors
curl --proto '=https' --tlsv1.2 --fail --location \
https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/resolve/f4a64596bbcf9a2a94591b74b9dc39b2e4e77e3e/vocab.txt \
--output models/chinese-clip/vocab.txt
printf '%s %s\n%s %s\n' \
29cc0b2bcf6ff777f2e15742be92b110e4acbdb2068356e862c4637a4b15fe4f \
models/chinese-clip/model.safetensors \
45bbac6b341c319adc98a532532882e91a9cefc0329aa57bac9ae761c27b291c \
models/chinese-clip/vocab.txt | sha256sum --check -
```
Compose 服务增加模型目录挂载与启动参数：
```yaml
services:
pixhelf:
volumes:
- "./pic:/data/pic:ro"
- pixhelf-cache:/data/.pixhelf-cache
- "./models/chinese-clip:/models/chinese-clip:ro"
command:
- "--text-search-model"
- "/models/chinese-clip/model.safetensors"
- "--text-search-vocab"
- "/models/chinese-clip/vocab.txt"
```
本地运行：
```bash
pixhelf \
--text-search-model ./models/chinese-clip/model.safetensors \
--text-search-vocab ./models/chinese-clip/vocab.txt
```
文字搜图使用一个 CPU 后台任务建立 512 维图像索引，向量按模型与词表指纹保存在缩略图缓存中，每张约 552 字节。侧边栏会显示“文字索引”进度；更换模型文件前请先停止 Pixhelf。
## 本地以图搜图
默认模式不需要模型：相似度完全在本机计算，综合比较图片的构图、亮度、色彩分布和边缘纹理，并过滤明显无关的结果。特征描述符会保存在缩略图缓存卷中；升级已有图库时会在后台自动补齐，之后重启无需重新计算。
如果希望识别“主体相近、构图和颜色不同”的图片，可以选择挂载 Meta 的 DINOv2 Small 官方权重。Pixhelf 不会联网或自动下载模型；开启后会用一个 CPU 后台任务建立 384 维语义索引，再与默认视觉特征混合排序。模型缺失或单张图片推理失败时会自动使用默认算法。
先在宿主机下载并校验固定版本的模型（88.2 MB，Apache-2.0）：
```bash
mkdir -p models
curl --proto '=https' --tlsv1.2 --fail --location \
https://huggingface.co/facebook/dinov2-small/resolve/ed25f3a31f01632728cabb09d1542f84ab7b0056/model.safetensors \
--output models/dinov2-small.safetensors
printf '%s %s\n' \
ae1e99fcefd534ed978cdeb8326f08030c96e28b7a81ffcbc98a857c84d14be1 \
models/dinov2-small.safetensors | sha256sum --check -
```
然后在 Compose 服务中增加只读挂载和启动参数：
```yaml
services:
pixhelf:
volumes:
- "./pic:/data/pic:ro"
- pixhelf-cache:/data/.pixhelf-cache
- "./models/dinov2-small.safetensors:/models/dinov2-small.safetensors:ro"
command: ["--semantic-model", "/models/dinov2-small.safetensors"]
```
本地运行时使用同一个参数：
```bash
pixhelf --semantic-model ./models/dinov2-small.safetensors
```
语义向量按模型指纹保存在缩略图缓存中，每张约 424 字节；侧边栏会显示“语义索引”进度。更换模型文件前请先停止 Pixhelf，替换后再启动。
frontend/scripts/visual-check.mjs
+26 -10
View File
@@ -328,16 +328,29 @@ try {
})(),
scrollCuePresentation: (() => {
const cue = document.querySelector(".viewer-scroll-cue");
const media = document.querySelector(".viewer-media");
if (!cue || !media) return null;
if (!cue) return null;
const style = getComputedStyle(cue);
const cueBounds = cue.getBoundingClientRect();
const mediaBounds = media.getBoundingClientRect();
return {
borderless: Number.parseFloat(style.borderTopWidth) === 0,
transparent: style.backgroundColor === "rgba(0, 0, 0, 0)",
surfaced: style.backgroundColor !== "rgba(0, 0, 0, 0)"
&& Number.parseFloat(style.borderTopWidth) > 0,
circular: Number.parseFloat(style.borderTopLeftRadius) >= cueBounds.height / 2 - 1,
width: cueBounds.width,
clearsMedia: cueBounds.top >= mediaBounds.bottom - 1,
};
})(),
floatingToolbar: (() => {
const viewer = document.querySelector(".image-viewer");
const header = document.querySelector(".viewer-header");
const actions = document.querySelector(".viewer-header-actions");
if (!viewer || !header || !actions) return null;
const headerStyle = getComputedStyle(header);
const actionsStyle = getComputedStyle(actions);
return {
mode: viewer.getAttribute("data-ui-layout"),
headerTransparent: headerStyle.backgroundImage === "none"
&& headerStyle.backgroundColor === "rgba(0, 0, 0, 0)",
actionsSurfaced: actionsStyle.backgroundColor !== "rgba(0, 0, 0, 0)"
&& Number.parseFloat(actionsStyle.borderTopWidth) > 0,
};
})(),
headingModules: document.querySelectorAll(".viewer-heading").length,
@@ -1628,10 +1641,13 @@ try {
viewer.switchPerformance.mediaAnimations !== 0 ||
viewer.switchPerformance.dispatchDuration > 80 ||
viewer.switchPerformance.firstFrameDuration > 100 ||
!viewerChrome.scrollCuePresentation?.borderless ||
!viewerChrome.scrollCuePresentation?.transparent ||
viewerChrome.scrollCuePresentation.width < 88 ||
!viewerChrome.scrollCuePresentation?.clearsMedia ||
!viewerChrome.scrollCuePresentation?.surfaced ||
!viewerChrome.scrollCuePresentation?.circular ||
viewerChrome.scrollCuePresentation.width < 44 ||
viewerChrome.scrollCuePresentation.width > 52 ||
viewerChrome.floatingToolbar?.mode !== "floating" ||
!viewerChrome.floatingToolbar?.headerTransparent ||
!viewerChrome.floatingToolbar?.actionsSurfaced ||
(target.name === "desktop" && (
!viewer.detailsReachedByWheel ||
viewer.detailsName !== viewer.keyboardNavigation ||
frontend/src/App.tsx
+310 -27
View File
@@ -23,7 +23,13 @@ import {
} from "preact/hooks";
import { ToolbarPopover } from "./ToolbarPopover";
import { ImageViewer } from "./ImageViewer";
import { getGallery, getImages, getStatus, takeInitialBootstrap } from "./api";
import {
getGallery,
getImages,
getSimilarImages,
getStatus,
takeInitialBootstrap,
} from "./api";
import type {
GalleryImage,
GallerySummary,
@@ -37,6 +43,7 @@ import {
} from "./viewerAssets";
const PAGE_SIZE = 60;
const SIMILAR_PAGE_SIZE = 30;
const CARD_PREFETCH_MARGIN = "1200px 0px";
const MOBILE_PAGE_PREFETCH_MARGIN = "1400px 0px";
const DESKTOP_PAGE_PREFETCH_MARGIN = "900px 0px";
@@ -57,6 +64,13 @@ type ImagePageState = {
nextOffset: number | null;
};
type SimilarImagePageState = {
source: GalleryImage | null;
images: GalleryImage[];
total: number;
nextOffset: number | null;
};
type ViewerAnchor = {
cardRatio: number;
viewportRatio: number;
@@ -85,6 +99,13 @@ const EMPTY_IMAGE_PAGE: ImagePageState = {
nextOffset: null,
};
const EMPTY_SIMILAR_IMAGE_PAGE: SimilarImagePageState = {
source: null,
images: [],
total: 0,
nextOffset: null,
};
function clampValue(value: number, minimum: number, maximum: number): number {
return Math.min(maximum, Math.max(minimum, value));
}
@@ -510,11 +531,34 @@ function App() {
const [mobileNavOpen, setMobileNavOpen] = useState(false);
const [sidebarCollapsed, setSidebarCollapsed] = useState(initialSidebarCollapsed);
const [viewerImageId, setViewerImageId] = useState<string | null>(null);
const [viewerDiscoveredImages, setViewerDiscoveredImages] = useState<GalleryImage[]>([]);
const [similarPage, setSimilarPage] = useState<SimilarImagePageState>(
EMPTY_SIMILAR_IMAGE_PAGE,
);
const [similarLoading, setSimilarLoading] = useState(false);
const [similarLoadingMore, setSimilarLoadingMore] = useState(false);
const [similarError, setSimilarError] = useState<string | null>(null);
const [viewerReturnRequest, setViewerReturnRequest] = useState<ViewerReturnRequest | null>(null);
const [reloadToken, setReloadToken] = useState(0);
const compactLayout = useMediaQuery("(max-width: 720px)");
const mobileNavMounted = useDelayedUnmount(mobileNavOpen, MOBILE_NAV_EXIT_MS);
const debouncedSearch = useDebounced(search.trim(), 250);
const textSearchReady = Boolean(
status?.backgroundComplete
&& status.textSearch.enabled
&& status.textSearch.backgroundComplete
&& status.textSearch.ready > 0,
);
const textSearchMode: "filename" | "indexing" | "semantic" = status?.textSearch.enabled
? textSearchReady
? "semantic"
: status.backgroundComplete && status.textSearch.backgroundComplete
? "filename"
: "indexing"
: "filename";
const debouncedSearch = useDebounced(
search.trim(),
status?.textSearch.enabled ? 500 : 250,
);
const requestVersionRef = useRef(0);
const summaryRevisionRef = useRef<string | null>(
INITIAL_BOOTSTRAP?.summary.revision ?? null,
@@ -522,11 +566,25 @@ function App() {
const skipInitialImagesRef = useRef(Boolean(INITIAL_BOOTSTRAP));
const loadMoreControllerRef = useRef<AbortController | null>(null);
const loadMorePromiseRef = useRef<Promise<GalleryImage[]> | null>(null);
const similarRequestVersionRef = useRef(0);
const similarControllerRef = useRef<AbortController | null>(null);
const similarLoadMoreControllerRef = useRef<AbortController | null>(null);
const similarLoadMorePromiseRef = useRef<Promise<GalleryImage[]> | null>(null);
const backgroundCompletionRef = useRef(
Boolean(
INITIAL_BOOTSTRAP?.status.backgroundComplete
&& INITIAL_BOOTSTRAP.status.semantic.backgroundComplete,
),
);
const viewerImageIdRef = useRef<string | null>(null);
const viewerAnchorRef = useRef<ViewerAnchor | null>(null);
const viewerReturnSequenceRef = useRef(0);
const viewerReturnFlightRef = useRef<ViewerReturnFlight | null>(null);
const { images, total, nextOffset } = imagePage;
const viewerImages = useMemo(
() => appendUniqueImages(images, viewerDiscoveredImages),
[images, viewerDiscoveredImages],
);
useLayoutEffect(() => {
window.scrollTo({ top: 0, left: 0, behavior: "auto" });
@@ -539,6 +597,8 @@ function App() {
useEffect(() => () => {
disposeViewerReturnFlight(viewerReturnFlightRef.current);
viewerReturnFlightRef.current = null;
similarControllerRef.current?.abort();
similarLoadMoreControllerRef.current?.abort();
}, []);
useEffect(() => {
@@ -626,13 +686,22 @@ function App() {
const next = await getStatus(controller.signal);
if (controller.signal.aborted) return;
setStatus(next);
timer = window.setTimeout(poll, next.backgroundComplete ? 10000 : 1500);
const complete = next.backgroundComplete
&& next.semantic.backgroundComplete
&& next.textSearch.backgroundComplete;
timer = window.setTimeout(poll, complete ? 10000 : 1500);
} catch {
if (!controller.signal.aborted) timer = window.setTimeout(poll, 5000);
}
};
const initialDelay = INITIAL_BOOTSTRAP
? (INITIAL_BOOTSTRAP.status.backgroundComplete ? 10_000 : 1_500)
? (
INITIAL_BOOTSTRAP.status.backgroundComplete
&& INITIAL_BOOTSTRAP.status.semantic.backgroundComplete
&& INITIAL_BOOTSTRAP.status.textSearch.backgroundComplete
? 10_000
: 1_500
)
: 0;
timer = window.setTimeout(poll, initialDelay);
return () => {
@@ -691,7 +760,7 @@ function App() {
loadMorePromiseRef.current = null;
loadMoreController?.abort();
};
}, [album, debouncedSearch, exploreSeed, reloadToken]);
}, [album, debouncedSearch, exploreSeed, reloadToken, textSearchReady]);
const loadMore = useCallback((): Promise<GalleryImage[]> => {
if (loading || nextOffset === null) return Promise.resolve([]);
@@ -740,6 +809,124 @@ function App() {
return promise;
}, [album, debouncedSearch, exploreSeed, loading, nextOffset]);
const resetSimilarSearch = useCallback(() => {
similarRequestVersionRef.current += 1;
similarControllerRef.current?.abort();
similarLoadMoreControllerRef.current?.abort();
similarControllerRef.current = null;
similarLoadMoreControllerRef.current = null;
similarLoadMorePromiseRef.current = null;
setSimilarPage(EMPTY_SIMILAR_IMAGE_PAGE);
setSimilarLoading(false);
setSimilarLoadingMore(false);
setSimilarError(null);
}, []);
const searchSimilar = useCallback((source: GalleryImage, force = false) => {
const sameSource = similarPage.source?.id === source.id;
if (!force && sameSource && (similarLoading || !similarError)) return;
similarControllerRef.current?.abort();
similarLoadMoreControllerRef.current?.abort();
similarLoadMoreControllerRef.current = null;
similarLoadMorePromiseRef.current = null;
const controller = new AbortController();
const requestVersion = ++similarRequestVersionRef.current;
similarControllerRef.current = controller;
setSimilarPage({ source, images: [], total: 0, nextOffset: null });
setSimilarLoading(true);
setSimilarLoadingMore(false);
setSimilarError(null);
void getSimilarImages(
source.id,
{ offset: 0, limit: SIMILAR_PAGE_SIZE },
controller.signal,
)
.then((page) => {
if (requestVersion !== similarRequestVersionRef.current) return;
setSimilarPage({
source,
images: page.items,
total: page.total,
nextOffset: page.nextOffset,
});
setViewerDiscoveredImages((current) => appendUniqueImages(current, page.items));
})
.catch((reason: unknown) => {
if (!controller.signal.aborted && requestVersion === similarRequestVersionRef.current) {
setSimilarError(errorMessage(reason, "无法查找相似图片"));
}
})
.finally(() => {
if (similarControllerRef.current === controller) {
similarControllerRef.current = null;
if (requestVersion === similarRequestVersionRef.current) setSimilarLoading(false);
}
});
}, [similarError, similarLoading, similarPage.source?.id]);
useEffect(() => {
const complete = Boolean(
status?.backgroundComplete && status.semantic.backgroundComplete,
);
const justCompleted = complete && !backgroundCompletionRef.current;
backgroundCompletionRef.current = complete;
if (justCompleted && similarPage.source) {
searchSimilar(similarPage.source, true);
}
}, [searchSimilar, similarPage.source, status]);
const loadMoreSimilar = useCallback((): Promise<GalleryImage[]> => {
const source = similarPage.source;
const offset = similarPage.nextOffset;
if (!source || similarLoading || offset === null) return Promise.resolve([]);
const pending = similarLoadMorePromiseRef.current;
if (pending) return pending;
const controller = new AbortController();
const requestVersion = similarRequestVersionRef.current;
similarLoadMoreControllerRef.current = controller;
setSimilarLoadingMore(true);
setSimilarError(null);
const promise = (async () => {
try {
const page = await getSimilarImages(
source.id,
{ offset, limit: SIMILAR_PAGE_SIZE },
controller.signal,
);
if (requestVersion !== similarRequestVersionRef.current) return [];
setSimilarPage((current) => current.source?.id === source.id
? {
source,
images: appendUniqueImages(current.images, page.items),
total: page.total,
nextOffset: page.nextOffset,
}
: current);
setViewerDiscoveredImages((current) => appendUniqueImages(current, page.items));
return page.items;
} catch (reason) {
if (!controller.signal.aborted && requestVersion === similarRequestVersionRef.current) {
setSimilarError(errorMessage(reason, "无法继续加载相似图片"));
}
return [];
} finally {
if (similarLoadMoreControllerRef.current === controller) {
similarLoadMoreControllerRef.current = null;
similarLoadMorePromiseRef.current = null;
if (requestVersion === similarRequestVersionRef.current) {
setSimilarLoadingMore(false);
}
}
}
})();
similarLoadMorePromiseRef.current = promise;
return promise;
}, [similarLoading, similarPage.nextOffset, similarPage.source]);
const sentinelRef = useRef<HTMLDivElement>(null);
useEffect(() => {
const sentinel = sentinelRef.current;
@@ -759,8 +946,22 @@ function App() {
}, [compactLayout, loadMore, nextOffset]);
const viewerIndex = viewerImageId
? viewerImages.findIndex((image) => image.id === viewerImageId)
: -1;
const galleryViewerIndex = viewerImageId
? images.findIndex((image) => image.id === viewerImageId)
: -1;
const viewerImage = viewerIndex >= 0 ? viewerImages[viewerIndex] : null;
const similarActive = Boolean(
viewerImage && similarPage.source?.id === viewerImage.id,
);
const viewerCanLoadMore = Boolean(
viewerImage && (
images.some((candidate) => candidate.id === viewerImage.id)
? nextOffset !== null
: similarActive && similarPage.nextOffset !== null
),
);
const openViewer = useCallback((
imageId: string,
@@ -771,10 +972,12 @@ function App() {
viewerReturnFlightRef.current = null;
viewerAnchorRef.current = captureViewerAnchor(card, pointerY);
prepareViewerImages(images, images.findIndex((image) => image.id === imageId));
resetSimilarSearch();
setViewerDiscoveredImages([]);
setViewerReturnRequest(null);
viewerImageIdRef.current = imageId;
setViewerImageId(imageId);
}, [images]);
}, [images, resetSimilarSearch]);
const closeViewer = useCallback(() => {
const currentImageId = viewerImageIdRef.current;
@@ -792,9 +995,11 @@ function App() {
imageId: currentImageId,
sequence: viewerReturnSequenceRef.current,
});
resetSimilarSearch();
setViewerDiscoveredImages([]);
viewerImageIdRef.current = null;
setViewerImageId(null);
}, []);
}, [resetSimilarSearch]);
useLayoutEffect(() => {
if (viewerImageId || !viewerReturnRequest) return;
@@ -1016,31 +1221,45 @@ function App() {
}, [viewerImageId, viewerReturnRequest]);
useEffect(() => {
if (!viewerImageId || viewerIndex < 0 || nextOffset === null) return;
if (viewerIndex >= images.length - 5) void loadMore();
}, [images.length, loadMore, nextOffset, viewerImageId, viewerIndex]);
if (!viewerImageId || galleryViewerIndex < 0 || nextOffset === null) return;
if (galleryViewerIndex >= images.length - 5) void loadMore();
}, [galleryViewerIndex, images.length, loadMore, nextOffset, viewerImageId]);
const openSimilarImage = useCallback((candidate: GalleryImage) => {
setViewerDiscoveredImages((current) => appendUniqueImages(current, [candidate]));
prepareViewerImages([candidate], 0);
viewerImageIdRef.current = candidate.id;
setViewerImageId(candidate.id);
searchSimilar(candidate);
}, [searchSimilar]);
const navigateViewer = useCallback((direction: -1 | 1) => {
const currentImageId = viewerImageIdRef.current;
if (!currentImageId) return;
const currentIndex = images.findIndex((image) => image.id === currentImageId);
const currentIndex = viewerImages.findIndex((image) => image.id === currentImageId);
if (currentIndex < 0) return;
const target = images[currentIndex + direction];
const target = viewerImages[currentIndex + direction];
if (target) {
viewerImageIdRef.current = target.id;
setViewerImageId(target.id);
return;
}
if (direction < 0 || nextOffset === null) return;
if (direction < 0) return;
const startingId = currentImageId;
void loadMore().then((incoming) => {
const pending = similarPage.source?.id === currentImageId
&& similarPage.nextOffset !== null
? loadMoreSimilar()
: nextOffset !== null && images.some((candidate) => candidate.id === currentImageId)
? loadMore()
: Promise.resolve([]);
void pending.then((incoming) => {
const next = incoming[0];
if (!next || viewerImageIdRef.current !== startingId) return;
viewerImageIdRef.current = next.id;
setViewerImageId(next.id);
});
}, [images, loadMore, nextOffset]);
}, [images, loadMore, loadMoreSimilar, nextOffset, similarPage, viewerImages]);
const galleryPath = `/${album.replace(/^\/+/, "")}`;
const chooseAlbum = (path: string) => {
@@ -1077,6 +1296,7 @@ function App() {
>
<Header
search={search}
searchMode={textSearchMode}
onSearchChange={changeSearch}
onExplore={startExploring}
exploreActive={Boolean(exploreSeed)}
@@ -1148,7 +1368,11 @@ function App() {
<ImageIcon size={30} strokeWidth={1.6} />
<strong>{summary?.total === 0 ? "图库暂无图片" : "没有找到图片"}</strong>
<span>
{summary?.total === 0 ? "添加图片后将自动显示" : "请调整相册或搜索条件"}
{summary?.total === 0
? "添加图片后将自动显示"
: textSearchReady && debouncedSearch
? "换一种自然语言描述，或减少限定词"
: "请调整相册或搜索条件"}
</span>
</div>
)}
@@ -1162,13 +1386,23 @@ function App() {
</main>
{viewerIndex >= 0 && (
<ImageViewer
images={images}
images={viewerImages}
activeIndex={viewerIndex}
total={total}
hasMore={nextOffset !== null}
loadingMore={loadingMore}
total={Math.max(summary?.total ?? total, viewerImages.length)}
hasMore={viewerCanLoadMore}
loadingMore={loadingMore || similarLoadingMore}
onNavigate={navigateViewer}
onClose={closeViewer}
similarActive={similarActive}
similarImages={similarActive ? similarPage.images : []}
similarTotal={similarActive ? similarPage.total : 0}
similarHasMore={similarActive && similarPage.nextOffset !== null}
similarLoading={similarActive && similarLoading}
similarLoadingMore={similarActive && similarLoadingMore}
similarError={similarActive ? similarError : null}
onSearchSimilar={searchSimilar}
onLoadMoreSimilar={() => void loadMoreSimilar()}
onOpenSimilar={openSimilarImage}
/>
)}
</div>
@@ -1177,6 +1411,7 @@ function App() {
function Header({
search,
searchMode,
onSearchChange,
onExplore,
exploreActive,
@@ -1186,6 +1421,7 @@ function Header({
navigationOpen,
}: {
search: string;
searchMode: "filename" | "indexing" | "semantic";
onSearchChange: (value: string) => void;
onExplore: () => void;
exploreActive: boolean;
@@ -1197,6 +1433,13 @@ function Header({
const [searchOpen, setSearchOpen] = useState(false);
const [exploreMotionKey, setExploreMotionKey] = useState(0);
const searchInputRef = useRef<HTMLInputElement>(null);
const semanticSearch = searchMode === "semantic";
const searchIndexing = searchMode === "indexing";
const searchPlaceholder = semanticSearch
? "描述想找的图片"
: searchIndexing
? "文字索引中 · 暂搜文件名"
: "搜索文件名";
useEffect(() => {
if (!searchOpen) return;
@@ -1280,7 +1523,7 @@ function Header({
<span className="toolbar-state-dot" aria-hidden="true" />
</>
}
rootClassName={`topbar-search ${search ? "has-query" : ""}`}
rootClassName={`topbar-search ${search ? "has-query" : ""} ${semanticSearch ? "is-semantic" : ""}`}
triggerClassName="search-toggle"
panelClassName="search-popover-panel"
>
@@ -1290,12 +1533,21 @@ function Header({
type="search"
value={search}
onInput={(event) => onSearchChange(event.currentTarget.value)}
placeholder="搜索文件名"
aria-label="搜索文件名"
placeholder={searchPlaceholder}
aria-label={semanticSearch ? "用自然语言搜索图片" : "搜索文件名"}
enterKeyHint="search"
autoComplete="off"
tabIndex={searchOpen ? 0 : -1}
/>
{searchMode !== "filename" && (
<span
className="search-mode-badge"
data-state={searchMode}
title={semanticSearch ? "Chinese-CLIP 本地语义搜索" : "正在建立文字搜图索引"}
>
{semanticSearch ? "语义" : "索引中"}
</span>
)}
{search && (
<button
type="button"
@@ -1352,16 +1604,47 @@ function SidebarStatus({ status }: { status: ThumbnailStatus | null }) {
</div>
);
}
const complete = status.backgroundComplete;
const thumbnailsComplete = status.backgroundComplete;
const semanticIndexing = status.semantic.enabled && !status.semantic.backgroundComplete;
const textSearchIndexing = status.textSearch.enabled && !status.textSearch.backgroundComplete;
const complete = thumbnailsComplete && !semanticIndexing && !textSearchIndexing;
const ready = !thumbnailsComplete
? status.ready
: textSearchIndexing
? status.textSearch.ready
: semanticIndexing
? status.semantic.ready
: status.ready;
const total = !thumbnailsComplete
? status.total
: textSearchIndexing
? status.textSearch.total
: semanticIndexing
? status.semantic.total
: status.total;
const label = !thumbnailsComplete
? "处理中"
: textSearchIndexing
? "文字索引"
: semanticIndexing
? "语义索引"
: "已就绪";
const title = !thumbnailsComplete
? "正在后台处理缩略图"
: textSearchIndexing
? "正在本地建立自然语言文字搜图索引"
: semanticIndexing
? "正在本地建立语义推荐索引"
: "图片和搜索索引处理完成";
return (
<div
className={`sidebar-status ${complete ? "complete" : ""}`}
title={complete ? "缩略图处理完成" : "正在后台处理缩略图"}
title={title}
role="status"
>
{complete ? <Check size={15} /> : <LoaderCircle className="spin" size={15} />}
<span>{complete ? "已就绪" : "处理中"}</span>
<strong>{formatCount(status.ready)} / {formatCount(status.total)}</strong>
<span>{label}</span>
<strong>{formatCount(ready)} / {formatCount(total)}</strong>
</div>
);
}
frontend/src/ImageViewer.tsx
+227 -3
View File
@@ -11,6 +11,7 @@ import {
Maximize2,
Minimize2,
RefreshCw,
ScanSearch,
X,
} from "./icons";
import type { GalleryImage } from "./types";
@@ -201,9 +202,8 @@ function viewerMediaDimensions(
image: GalleryImage,
viewport: ViewportSize,
): { width: number; height: number } {
const compact = viewport.width <= 720;
const availableWidth = Math.max(1, viewport.width - (compact ? 12 : 144));
const availableHeight = Math.max(1, viewport.height - (compact ? 128 : 104));
const availableWidth = Math.max(1, viewport.width);
const availableHeight = Math.max(1, viewport.height);
const imageRatio = image.width / Math.max(1, image.height);
const width = Math.min(availableWidth, availableHeight * imageRatio);
return { width, height: width / imageRatio };
@@ -297,6 +297,131 @@ function drawViewerCanvas(
canvas.dataset.renderedSource = sourceKey;
}
function SimilarImageCard({
image,
onOpen,
}: {
image: GalleryImage;
onOpen: (image: GalleryImage) => void;
}) {
return (
<figure
className="viewer-similar-card"
title={image.name}
data-image-id={image.id}
role="button"
tabIndex={0}
aria-label={`查看相似图片 ${image.name}`}
onClick={() => onOpen(image)}
onKeyDown={(event) => {
if (event.key !== "Enter" && event.key !== " ") return;
event.preventDefault();
onOpen(image);
}}
>
<div className="viewer-similar-card-media">
<img
src={viewerThumbnailUrl(image)}
alt={image.name}
loading="lazy"
decoding="async"
width={image.width}
height={image.height}
draggable={false}
onError={(event) => {
event.currentTarget.dataset.failed = "true";
}}
/>
<span className="viewer-similar-card-fallback" aria-hidden="true">
<RefreshCw size={18} />
</span>
</div>
<figcaption>{image.name}</figcaption>
</figure>
);
}
function SimilarImageMasonry({
images,
total,
columnCount,
hasMore,
loadingMore,
error,
onLoadMore,
onOpen,
}: {
images: GalleryImage[];
total: number;
columnCount: number;
hasMore: boolean;
loadingMore: boolean;
error: string | null;
onLoadMore: () => void;
onOpen: (image: GalleryImage) => void;
}) {
const columns = Array.from({ length: columnCount }, () => [] as GalleryImage[]);
const heights = Array.from({ length: columnCount }, () => 0);
for (const image of images) {
const target = heights.indexOf(Math.min(...heights));
columns[target]!.push(image);
heights[target] += image.height / Math.max(1, image.width) + 0.12;
}
return (
<>
<div
className="viewer-similar-masonry"
data-columns={columnCount}
style={{ "--similar-columns": columnCount } as CSSProperties}
>
{columns.map((column, index) => (
<div key={index} className="viewer-similar-column">
{column.map((candidate) => (
<SimilarImageCard key={candidate.id} image={candidate} onOpen={onOpen} />
))}
</div>
))}
</div>
{error && (
<div className="viewer-similar-inline-error" role="alert">
<span>{error}</span>
<button type="button" onClick={onLoadMore}>重试</button>
</div>
)}
{!error && (
<div className="viewer-similar-pagination">
{hasMore ? (
<button
type="button"
className="viewer-similar-load-more"
onClick={onLoadMore}
disabled={loadingMore}
aria-busy={loadingMore}
>
{loadingMore && <LoaderCircle className="spin" size={15} />}
<span>{loadingMore ? "正在加载" : "加载更多相似图片"}</span>
{!loadingMore && <small>{images.length} / {total}</small>}
</button>
) : (
<span className="viewer-similar-count">
已显示全部 {total.toLocaleString()} 张相似图片
</span>
)}
</div>
)}
</>
);
}
function SimilarImageSkeleton() {
return (
<div className="viewer-similar-skeleton" aria-label="正在查找相似图片">
{[1, 2, 3, 4, 5, 6].map((item) => <span key={item} />)}
</div>
);
}
export function ImageViewer({
images,
activeIndex,
@@ -305,6 +430,16 @@ export function ImageViewer({
loadingMore,
onNavigate,
onClose,
similarActive,
similarImages,
similarTotal,
similarHasMore,
similarLoading,
similarLoadingMore,
similarError,
onSearchSimilar,
onLoadMoreSimilar,
onOpenSimilar,
}: {
images: GalleryImage[];
activeIndex: number;
@@ -313,6 +448,16 @@ export function ImageViewer({
loadingMore: boolean;
onNavigate: (direction: -1 | 1) => void;
onClose: () => void;
similarActive: boolean;
similarImages: GalleryImage[];
similarTotal: number;
similarHasMore: boolean;
similarLoading: boolean;
similarLoadingMore: boolean;
similarError: string | null;
onSearchSimilar: (image: GalleryImage) => void;
onLoadMoreSimilar: () => void;
onOpenSimilar: (image: GalleryImage) => void;
}) {
const image = images[activeIndex]!;
const viewport = useViewportSize();
@@ -333,6 +478,7 @@ export function ImageViewer({
const waitingForNext = loadingMore && activeIndex === images.length - 1;
const dialogRef = useRef<HTMLDivElement>(null);
const detailsSectionRef = useRef<HTMLElement>(null);
const similarSectionRef = useRef<HTMLElement>(null);
const closeButtonRef = useRef<HTMLButtonElement>(null);
const surfaceRef = useRef<HTMLDivElement>(null);
const mediaRef = useRef<HTMLDivElement>(null);
@@ -613,6 +759,17 @@ export function ImageViewer({
}
};
const scrollToSimilar = () => {
commitTransform(DEFAULT_TRANSFORM);
wheelZoomBlockedUntilRef.current = performance.now() + WHEEL_HANDOFF_DELAY_MS;
window.requestAnimationFrame(() => {
similarSectionRef.current?.scrollIntoView({
behavior: reducedMotion() ? "auto" : "smooth",
block: "start",
});
});
};
const parkViewerControlFocus = () => {
const viewer = dialogRef.current;
const active = document.activeElement;
@@ -1775,6 +1932,7 @@ export function ImageViewer({
data-page={scrollPage}
data-renderer={renderer}
data-control-system="unified"
data-ui-layout="floating"
data-scroll-mode="continuous"
data-mouse-side-navigation="true"
data-mobile-swipe-motion="interruptible"
@@ -1789,6 +1947,19 @@ export function ImageViewer({
<header className="viewer-header">
<div className="viewer-header-actions">
<button
type="button"
className={`viewer-control viewer-similar-trigger${similarActive ? " is-active" : ""}`}
onClick={() => {
onSearchSimilar(image);
scrollToSimilar();
}}
aria-label="以图搜图，查找相似图片"
aria-pressed={similarActive}
title="以图搜图"
>
<ScanSearch size={19} />
</button>
<a
className="viewer-control"
href={viewerOriginalUrl(image)}
@@ -2062,6 +2233,59 @@ export function ImageViewer({
</span>
<ChevronRight size={18} />
</a>
<section
ref={similarSectionRef}
className="viewer-similar-section"
aria-labelledby="viewer-similar-title"
data-similar-active={similarActive}
>
<div className="viewer-similar-heading">
<div>
<h3 id="viewer-similar-title">相似图片</h3>
<p>在图库中查找构图、色彩和纹理相近的图片</p>
</div>
<ScanSearch size={20} aria-hidden="true" />
</div>
{!similarActive ? (
<button
type="button"
className="viewer-similar-search-button"
onClick={() => onSearchSimilar(image)}
>
<ScanSearch size={18} />
<span>以图搜图</span>
<ChevronRight size={17} />
</button>
) : similarLoading && !similarImages.length ? (
<SimilarImageSkeleton />
) : similarError && !similarImages.length ? (
<div className="viewer-similar-error" role="alert">
<span>{similarError}</span>
<button type="button" onClick={() => onSearchSimilar(image)}>
<RefreshCw size={15} />
重试
</button>
</div>
) : similarImages.length ? (
<SimilarImageMasonry
images={similarImages}
total={similarTotal}
columnCount={compactViewport ? 2 : 3}
hasMore={similarHasMore}
loadingMore={similarLoadingMore}
error={similarError}
onLoadMore={onLoadMoreSimilar}
onOpen={onOpenSimilar}
/>
) : (
<div className="viewer-similar-empty">
<ScanSearch size={22} />
<span>暂时没有找到相似图片</span>
</div>
)}
</section>
</div>
</div>
</section>
frontend/src/api.ts
+44 -3
View File
@@ -9,6 +9,8 @@ import type {
} from "./types";
const REQUEST_TIMEOUT_MS = 15_000;
const IMAGE_SEARCH_REQUEST_TIMEOUT_MS = 60_000;
const SIMILAR_REQUEST_TIMEOUT_MS = 60_000;
type JsonObject = Record<string, unknown>;
type JsonValidator<T> = (value: unknown) => value is T;
@@ -66,7 +68,23 @@ function isThumbnailStatus(value: unknown): value is ThumbnailStatus {
&& isNonNegativeInteger(value.processing)
&& isNonNegativeInteger(value.failed)
&& typeof value.initialBatchReady === "boolean"
&& typeof value.backgroundComplete === "boolean";
&& typeof value.backgroundComplete === "boolean"
&& isObject(value.semantic)
&& typeof value.semantic.enabled === "boolean"
&& isNonNegativeInteger(value.semantic.total)
&& isNonNegativeInteger(value.semantic.ready)
&& isNonNegativeInteger(value.semantic.queued)
&& isNonNegativeInteger(value.semantic.processing)
&& isNonNegativeInteger(value.semantic.failed)
&& typeof value.semantic.backgroundComplete === "boolean"
&& isObject(value.textSearch)
&& typeof value.textSearch.enabled === "boolean"
&& isNonNegativeInteger(value.textSearch.total)
&& isNonNegativeInteger(value.textSearch.ready)
&& isNonNegativeInteger(value.textSearch.queued)
&& isNonNegativeInteger(value.textSearch.processing)
&& isNonNegativeInteger(value.textSearch.failed)
&& typeof value.textSearch.backgroundComplete === "boolean";
}
function isBootstrapData(value: unknown): value is BootstrapData {
@@ -80,6 +98,7 @@ async function getJson<T>(
url: string,
validate: JsonValidator<T>,
signal?: AbortSignal,
timeoutMs = REQUEST_TIMEOUT_MS,
): Promise<T> {
const controller = new AbortController();
let timedOut = false;
@@ -93,7 +112,7 @@ async function getJson<T>(
const timeout = window.setTimeout(() => {
timedOut = true;
controller.abort();
}, REQUEST_TIMEOUT_MS);
}, timeoutMs);
try {
const response = await fetch(url, {
@@ -168,5 +187,27 @@ export function getImages(
if (options.album) params.set("album", options.album);
if (options.search) params.set("search", options.search);
if (options.seed) params.set("seed", options.seed);
return getJson(`/api/images?${params}`, isImagesPage, signal);
return getJson(
`/api/images?${params}`,
isImagesPage,
signal,
options.search ? IMAGE_SEARCH_REQUEST_TIMEOUT_MS : REQUEST_TIMEOUT_MS,
);
}
export function getSimilarImages(
imageId: string,
options: { offset: number; limit: number },
signal?: AbortSignal,
): Promise<ImagesPage> {
const params = new URLSearchParams({
offset: String(options.offset),
limit: String(options.limit),
});
return getJson(
`/api/images/${encodeURIComponent(imageId)}/similar?${params}`,
isImagesPage,
signal,
SIMILAR_REQUEST_TIMEOUT_MS,
);
}
frontend/src/icons.tsx
+13
View File
@@ -170,6 +170,19 @@ export function Search(props: IconProps) {
return <Icon {...props}><path d="m21 21-4.34-4.34" /><circle cx="11" cy="11" r="8" /></Icon>;
}
export function ScanSearch(props: IconProps) {
return (
<Icon {...props}>
<path d="M3 7V5a2 2 0 0 1 2-2h2" />
<path d="M17 3h2a2 2 0 0 1 2 2v2" />
<path d="M21 17v2a2 2 0 0 1-2 2h-2" />
<path d="M7 21H5a2 2 0 0 1-2-2v-2" />
<circle cx="10.8" cy="10.8" r="3.8" />
<path d="m14 14 3.2 3.2" />
</Icon>
);
}
export function X(props: IconProps) {
return <Icon {...props}><path d="M18 6 6 18" /><path d="m6 6 12 12" /></Icon>;
}
frontend/src/styles.css
+383 -29
View File
@@ -643,6 +643,10 @@ select:focus-visible,
--popover-width: 320px;
}
.topbar-search.is-semantic {
--popover-width: 356px;
}
.search-popover-panel {
height: 46px;
display: flex;
@@ -684,6 +688,25 @@ select:focus-visible,
.search-popover-panel input[type="search"]::-webkit-search-cancel-button { display: none; }
.search-mode-badge {
flex: 0 0 auto;
padding: 3px 6px;
color: #176c63;
background: #e8f3f1;
border: 1px solid #c8dfda;
border-radius: 999px;
font-size: 9px;
font-weight: 680;
line-height: 1.2;
white-space: nowrap;
}
.search-mode-badge[data-state="indexing"] {
color: #747b78;
background: #f1f2f1;
border-color: #e0e3e1;
}
.clear-search {
width: 25px;
height: 25px;
@@ -979,7 +1002,6 @@ select:focus-visible,
isolation: isolate;
outline: none;
overscroll-behavior: contain;
scrollbar-gutter: stable;
scrollbar-color: rgba(255, 255, 255, .18) transparent;
scrollbar-width: thin;
scroll-behavior: smooth;
@@ -1013,12 +1035,10 @@ select:focus-visible,
position: absolute;
z-index: 5;
inset: 0 0 auto;
min-height: 76px;
display: flex;
align-items: flex-start;
justify-content: flex-end;
padding: max(16px, env(safe-area-inset-top)) max(18px, env(safe-area-inset-right)) 18px max(18px, env(safe-area-inset-left));
background: linear-gradient(to bottom, rgba(5, 7, 6, .72), rgba(5, 7, 6, 0));
padding: max(14px, env(safe-area-inset-top)) max(14px, env(safe-area-inset-right)) 0 max(14px, env(safe-area-inset-left));
pointer-events: none;
transition: opacity var(--motion-fast) ease;
}
@@ -1026,10 +1046,33 @@ select:focus-visible,
.viewer-header-actions {
display: flex;
align-items: center;
gap: 7px;
gap: 3px;
padding: 4px;
background: rgba(10, 13, 12, .58);
border: 1px solid rgba(255, 255, 255, .13);
border-radius: 999px;
box-shadow:
0 12px 38px rgba(0, 0, 0, .32),
inset 0 1px rgba(255, 255, 255, .045);
backdrop-filter: blur(18px) saturate(1.14);
-webkit-backdrop-filter: blur(18px) saturate(1.14);
pointer-events: auto;
}
.viewer-header .viewer-control {
background: transparent;
border-color: transparent;
box-shadow: none;
backdrop-filter: none;
-webkit-backdrop-filter: none;
}
.viewer-header .viewer-control:hover,
.viewer-header .viewer-control:focus-visible {
background: rgba(255, 255, 255, .105);
border-color: transparent;
}
.viewer-control,
.viewer-nav,
.viewer-scroll-cue {
@@ -1217,7 +1260,7 @@ select:focus-visible,
position: absolute;
z-index: 5;
right: 0;
bottom: max(14px, env(safe-area-inset-bottom));
bottom: max(16px, env(safe-area-inset-bottom));
left: 0;
display: flex;
flex-direction: column;
@@ -1228,16 +1271,18 @@ select:focus-visible,
.viewer-scroll-cue {
position: relative;
width: 92px;
height: 34px;
width: 50px;
height: 50px;
padding: 0;
color: rgba(248, 250, 249, .76);
background: transparent;
border: 0;
border-radius: 0;
box-shadow: none;
backdrop-filter: none;
-webkit-backdrop-filter: none;
background: rgba(18, 22, 20, .6);
border: 1px solid rgba(255, 255, 255, .13);
border-radius: 50%;
box-shadow:
0 10px 32px rgba(0, 0, 0, .3),
inset 0 1px rgba(255, 255, 255, .04);
backdrop-filter: blur(16px) saturate(1.12);
-webkit-backdrop-filter: blur(16px) saturate(1.12);
cursor: pointer;
pointer-events: auto;
transition:
@@ -1249,10 +1294,9 @@ select:focus-visible,
.viewer-scroll-cue:hover,
.viewer-scroll-cue:focus-visible {
color: #fff;
background: transparent;
border-color: transparent;
background: rgba(48, 55, 52, .8);
border-color: rgba(255, 255, 255, .23);
outline: none;
filter: drop-shadow(0 3px 8px rgba(0, 0, 0, .58));
transform: translateY(2px);
}
@@ -1490,6 +1534,301 @@ select:focus-visible,
font-size: 10px;
}
.viewer-similar-trigger.is-active {
color: #b8eee7;
background: rgba(40, 122, 114, .48);
border-color: rgba(143, 210, 201, .38);
}
.viewer-similar-section {
scroll-margin-top: 90px;
display: flex;
flex-direction: column;
gap: 16px;
padding-top: 2px;
}
.viewer-similar-heading {
min-width: 0;
display: flex;
align-items: center;
justify-content: space-between;
gap: 18px;
padding: 0 2px;
}
.viewer-similar-heading > div {
min-width: 0;
display: flex;
flex-direction: column;
gap: 4px;
}
.viewer-similar-heading h3,
.viewer-similar-heading p {
margin: 0;
}
.viewer-similar-heading h3 {
color: rgba(250, 252, 251, .94);
font-size: 14px;
font-weight: 680;
}
.viewer-similar-heading p {
color: rgba(239, 243, 241, .46);
font-size: 10px;
}
.viewer-similar-heading > svg {
flex: 0 0 auto;
color: rgba(143, 210, 201, .64);
}
.viewer-similar-search-button,
.viewer-similar-load-more {
color: rgba(247, 250, 249, .86);
background: rgba(255, 255, 255, .05);
border: 1px solid rgba(255, 255, 255, .085);
transition:
color var(--motion-fast) ease,
background var(--motion-fast) ease,
border-color var(--motion-fast) ease,
transform var(--motion-fast) var(--ease-standard);
}
.viewer-similar-search-button {
min-height: 62px;
display: grid;
grid-template-columns: 24px minmax(0, 1fr) 20px;
align-items: center;
gap: 13px;
padding: 11px 16px;
text-align: left;
border-radius: 15px;
}
.viewer-similar-search-button:hover,
.viewer-similar-load-more:hover:not(:disabled) {
color: #fff;
background: rgba(74, 135, 126, .16);
border-color: rgba(143, 210, 201, .2);
}
.viewer-similar-search-button:active,
.viewer-similar-load-more:active:not(:disabled) {
transform: scale(.988);
}
.viewer-similar-search-button > span {
font-size: 12px;
font-weight: 650;
}
.viewer-similar-masonry {
width: 100%;
display: grid;
grid-template-columns: repeat(var(--similar-columns), minmax(0, 1fr));
align-items: start;
gap: 10px;
}
.viewer-similar-column {
min-width: 0;
display: flex;
flex-direction: column;
gap: 10px;
}
.viewer-similar-card {
width: 100%;
display: block;
overflow: hidden;
margin: 0;
color: rgba(247, 250, 249, .82);
background: rgba(255, 255, 255, .045);
border: 1px solid rgba(255, 255, 255, .075);
border-radius: 10px;
box-shadow: 0 8px 28px rgba(0, 0, 0, .12);
cursor: zoom-in;
transition:
background var(--motion-fast) ease,
border-color var(--motion-fast) ease,
box-shadow var(--motion-medium) var(--ease-standard),
transform var(--motion-medium) var(--ease-standard);
}
.viewer-similar-card:hover,
.viewer-similar-card:focus-visible {
color: #fff;
background: rgba(255, 255, 255, .075);
border-color: rgba(143, 210, 201, .24);
box-shadow: 0 12px 34px rgba(0, 0, 0, .2);
transform: translateY(-2px);
}
.viewer-similar-card:focus-visible {
outline: 2px solid #8fd2c9;
outline-offset: 3px;
}
.viewer-similar-card-media {
position: relative;
min-height: 80px;
overflow: hidden;
background: #090c0b;
}
.viewer-similar-card-media img {
position: relative;
z-index: 1;
width: 100%;
height: auto;
display: block;
object-fit: cover;
transition: opacity var(--motion-fast) ease, transform var(--motion-slow) var(--ease-standard);
}
.viewer-similar-card:hover .viewer-similar-card-media img {
transform: scale(1.018);
}
.viewer-similar-card-media img[data-failed="true"] {
opacity: 0;
}
.viewer-similar-card-fallback {
position: absolute;
inset: 0;
display: grid;
place-items: center;
color: rgba(239, 243, 241, .34);
}
.viewer-similar-card figcaption {
overflow: hidden;
padding: 9px 10px 10px;
font-size: 10px;
line-height: 1.35;
text-overflow: ellipsis;
white-space: nowrap;
}
.viewer-similar-pagination {
min-height: 48px;
display: flex;
align-items: center;
justify-content: center;
}
.viewer-similar-load-more {
min-width: 210px;
min-height: 40px;
display: flex;
align-items: center;
justify-content: center;
gap: 8px;
padding: 9px 15px;
border-radius: 999px;
font-size: 11px;
font-weight: 620;
}
.viewer-similar-load-more:disabled {
opacity: .56;
cursor: wait;
}
.viewer-similar-load-more small {
padding-left: 8px;
color: rgba(239, 243, 241, .42);
border-left: 1px solid rgba(255, 255, 255, .12);
font-size: 9px;
font-variant-numeric: tabular-nums;
font-weight: 500;
}
.viewer-similar-count {
color: rgba(239, 243, 241, .42);
font-size: 10px;
}
.viewer-similar-skeleton {
display: grid;
grid-template-columns: repeat(3, minmax(0, 1fr));
align-items: start;
gap: 10px;
}
.viewer-similar-skeleton span {
min-height: 120px;
background: linear-gradient(105deg, #242927 24%, #303734 43%, #242927 62%);
background-size: 220% 100%;
border-radius: 10px;
animation: skeleton-shimmer 1.35s ease-in-out infinite;
}
.viewer-similar-skeleton span:nth-child(2),
.viewer-similar-skeleton span:nth-child(5) { aspect-ratio: .76; }
.viewer-similar-skeleton span:nth-child(3),
.viewer-similar-skeleton span:nth-child(4) { aspect-ratio: 1.28; }
.viewer-similar-error,
.viewer-similar-empty {
min-height: 112px;
display: flex;
align-items: center;
justify-content: center;
gap: 10px;
padding: 20px;
color: rgba(239, 243, 241, .54);
background: rgba(255, 255, 255, .035);
border: 1px solid rgba(255, 255, 255, .065);
border-radius: 13px;
font-size: 11px;
}
.viewer-similar-error {
color: #ffd9d4;
}
.viewer-similar-error button {
min-height: 30px;
display: flex;
align-items: center;
gap: 6px;
padding: 5px 9px;
color: inherit;
background: rgba(124, 58, 49, .32);
border: 1px solid rgba(255, 187, 177, .16);
border-radius: 999px;
font-size: 10px;
}
.viewer-similar-inline-error {
min-height: 38px;
display: flex;
align-items: center;
justify-content: center;
gap: 9px;
padding: 7px 10px;
color: #ffd9d4;
background: rgba(124, 58, 49, .12);
border: 1px solid rgba(255, 187, 177, .12);
border-radius: 9px;
font-size: 10px;
}
.viewer-similar-inline-error button {
padding: 4px 8px;
color: inherit;
background: rgba(124, 58, 49, .28);
border: 1px solid rgba(255, 187, 177, .15);
border-radius: 999px;
font-size: 9px;
}
.viewer-load-error {
position: absolute;
z-index: 5;
@@ -1658,14 +1997,13 @@ select:focus-visible,
transform: translateY(0);
}
.viewer-header {
min-height: calc(64px + env(safe-area-inset-top));
padding:
calc(10px + env(safe-area-inset-top))
max(10px, env(safe-area-inset-right))
16px
max(12px, env(safe-area-inset-left));
max(8px, env(safe-area-inset-top))
max(8px, env(safe-area-inset-right))
0
max(8px, env(safe-area-inset-left));
}
.viewer-header-actions { gap: 6px; }
.viewer-header-actions { gap: 2px; padding: 3px; }
.viewer-control {
width: 40px;
height: 40px;
@@ -1682,11 +2020,11 @@ select:focus-visible,
}
.viewer-nav { display: none; }
.viewer-bottom-bar {
bottom: max(4px, env(safe-area-inset-bottom));
bottom: max(10px, env(safe-area-inset-bottom));
}
.viewer-scroll-cue {
width: 96px;
height: 30px;
width: 48px;
height: 48px;
}
.viewer-details-inner { width: calc(100% - 24px); }
.viewer-details-header {
@@ -1711,6 +2049,22 @@ select:focus-visible,
.viewer-details-summary h3 { font-size: 15px; }
.viewer-details-grid { grid-template-columns: repeat(2, minmax(0, 1fr)); }
.viewer-details-grid > div { padding: 13px 14px; }
.viewer-similar-section { scroll-margin-top: 76px; }
.viewer-similar-heading p { max-width: 240px; }
.viewer-similar-masonry {
gap: 6px;
}
.viewer-similar-column { gap: 6px; }
.viewer-similar-card {
border-radius: 8px;
}
.viewer-similar-card figcaption { padding: 8px; font-size: 9px; }
.viewer-similar-skeleton {
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 6px;
}
.viewer-similar-skeleton span:nth-child(n + 5) { display: none; }
.viewer-similar-load-more { width: 100%; }
.viewer-load-error {
bottom: max(64px, calc(52px + env(safe-area-inset-bottom)));
}
@@ -1762,10 +2116,10 @@ select:focus-visible,
}
@media (max-height: 520px) {
.viewer-bottom-bar { bottom: max(6px, env(safe-area-inset-bottom)); }
.viewer-bottom-bar { bottom: max(8px, env(safe-area-inset-bottom)); }
.viewer-scroll-cue {
width: 88px;
height: 28px;
width: 44px;
height: 44px;
}
.viewer-load-error { bottom: max(58px, calc(48px + env(safe-area-inset-bottom))); }
}
frontend/src/types.ts
+18
View File
@@ -33,6 +33,24 @@ export interface ThumbnailStatus {
failed: number;
initialBatchReady: boolean;
backgroundComplete: boolean;
semantic: {
enabled: boolean;
total: number;
ready: number;
queued: number;
processing: number;
failed: number;
backgroundComplete: boolean;
};
textSearch: {
enabled: boolean;
total: number;
ready: number;
queued: number;
processing: number;
failed: number;
backgroundComplete: boolean;
};
}
export interface BootstrapData {
src/config.rs
+63
View File
@@ -6,17 +6,28 @@ use anyhow::{Context, Result, bail};
pub struct Config {
pub gallery_dir: PathBuf,
pub cache_dir: PathBuf,
pub semantic_model: Option<PathBuf>,
pub text_search: Option<TextSearchFiles>,
pub listen: SocketAddr,
pub initial_batch: usize,
pub workers: usize,
pub scan_interval: Duration,
}
#[derive(Clone, Debug)]
pub struct TextSearchFiles {
pub model: PathBuf,
pub vocabulary: PathBuf,
}
impl Config {
pub fn from_env_and_args() -> Result<Self> {
let cwd = env::current_dir().context("cannot determine the working directory")?;
let mut gallery_dir = cwd.join("pic");
let mut cache_dir = cwd.join(".pixhelf-cache").join("thumbnails");
let mut semantic_model = None;
let mut text_search_model = None;
let mut text_search_vocabulary = None;
let mut listen = "0.0.0.0:3002".parse::<SocketAddr>()?;
let mut initial_batch = 60usize;
let mut workers = default_worker_count();
@@ -27,6 +38,15 @@ impl Config {
match arg.as_str() {
"--gallery-dir" => gallery_dir = path_value(&mut args, &arg, &cwd)?,
"--cache-dir" => cache_dir = path_value(&mut args, &arg, &cwd)?,
"--semantic-model" => {
semantic_model = Some(path_value(&mut args, &arg, &cwd)?);
}
"--text-search-model" => {
text_search_model = Some(path_value(&mut args, &arg, &cwd)?);
}
"--text-search-vocab" => {
text_search_vocabulary = Some(path_value(&mut args, &arg, &cwd)?);
}
"--listen" => {
listen = value(&mut args, &arg)?
.parse()
@@ -53,10 +73,28 @@ impl Config {
.with_context(|| format!("cannot create cache directory: {}", cache_dir.display()))?;
cache_dir = canonical_directory(&cache_dir, "cache")?;
ensure_cache_is_external(&gallery_dir, &cache_dir)?;
semantic_model = semantic_model
.map(|path| canonical_file(&path, "semantic model"))
.transpose()?;
let text_search = match text_search_model {
Some(model) => {
let model = canonical_file(&model, "text-search model")?;
let vocabulary = text_search_vocabulary
.unwrap_or_else(|| model.parent().unwrap_or(&cwd).join("vocab.txt"));
let vocabulary = canonical_file(&vocabulary, "text-search vocabulary")?;
Some(TextSearchFiles { model, vocabulary })
}
None if text_search_vocabulary.is_some() => {
bail!("--text-search-vocab requires --text-search-model");
}
None => None,
};
Ok(Self {
gallery_dir,
cache_dir,
semantic_model,
text_search,
listen,
initial_batch,
workers,
@@ -73,6 +111,14 @@ fn canonical_directory(path: &std::path::Path, label: &str) -> Result<PathBuf> {
.with_context(|| format!("cannot resolve {label} directory: {}", path.display()))
}
fn canonical_file(path: &std::path::Path, label: &str) -> Result<PathBuf> {
if !path.is_file() {
bail!("{label} file does not exist: {}", path.display());
}
std::fs::canonicalize(path)
.with_context(|| format!("cannot resolve {label} file: {}", path.display()))
}
fn ensure_cache_is_external(
gallery_dir: &std::path::Path,
cache_dir: &std::path::Path,
@@ -133,6 +179,9 @@ Usage: pixhelf [OPTIONS]\n\n\
Options:\n \
--gallery-dir PATH Gallery root (default: ./pic)\n \
--cache-dir PATH Thumbnail cache (default: ./.pixhelf-cache/thumbnails)\n \
--semantic-model PATH Optional official facebook/dinov2-small safetensors model\n \
--text-search-model PATH Optional OFA Chinese-CLIP ViT-B/16 safetensors model\n \
--text-search-vocab PATH Chinese-CLIP vocab.txt (default: beside model)\n \
--listen HOST:PORT Listen address (default: 0.0.0.0:3002)\n \
--initial-batch N Priority thumbnail batch (default: 60)\n \
--workers N Background thumbnail workers (default: 2-4, based on CPU)\n \
@@ -164,6 +213,20 @@ mod tests {
assert!(canonical_directory(&temp.path().join("missing"), "test").is_err());
}
#[test]
fn resolves_only_existing_model_files() {
let temp = tempfile::tempdir().unwrap();
let model = temp.path().join("model.safetensors");
std::fs::write(&model, b"model").unwrap();
assert_eq!(
canonical_file(&model, "model").unwrap(),
model.canonicalize().unwrap()
);
assert!(canonical_file(temp.path(), "model").is_err());
assert!(canonical_file(&temp.path().join("missing"), "model").is_err());
}
#[test]
fn rejects_cache_directories_inside_the_gallery() {
let temp = tempfile::tempdir().unwrap();
src/main.rs
+43 -1
View File
@@ -1,5 +1,8 @@
mod config;
mod gallery;
mod semantic;
mod similarity;
mod text_search;
mod thumbs;
mod web;
@@ -8,6 +11,8 @@ use std::{path::PathBuf, sync::Arc};
use anyhow::{Context, Result};
use config::Config;
use gallery::{GalleryIndex, scan_gallery};
use semantic::SemanticIndex;
use text_search::TextSearchIndex;
use thumbs::ThumbnailManager;
use tokio::{net::TcpListener, sync::RwLock};
use tracing::{error, info, warn};
@@ -46,12 +51,48 @@ async fn main() -> Result<()> {
.map(|record| record.id.clone())
.collect();
let index = Arc::new(RwLock::new(Arc::new(index)));
let thumbnails = ThumbnailManager::new(config.cache_dir.clone())?;
let semantic = match config.semantic_model.as_deref() {
Some(model_path) => match SemanticIndex::load(model_path) {
Ok(index) => Some(index),
Err(error) => {
warn!(path = %model_path.display(), %error, "semantic model unavailable; using local visual fallback");
None
}
},
None => None,
};
let text_search = match config.text_search.as_ref() {
Some(files) => match TextSearchIndex::load(&files.model, &files.vocabulary) {
Ok(index) => Some(index),
Err(error) => {
warn!(
model = %files.model.display(),
vocabulary = %files.vocabulary.display(),
%error,
"natural-language search model unavailable; using filename search"
);
None
}
},
None => None,
};
let thumbnails = ThumbnailManager::new_with_indexes(
config.cache_dir.clone(),
semantic.as_ref().map(Arc::clone),
text_search.as_ref().map(Arc::clone),
)?;
{
let current = index.read().await;
thumbnails.reconcile(&current.images);
}
thumbnails.start_workers(config.workers);
if let Some(semantic) = semantic {
semantic.start_worker();
}
if let Some(text_search) = text_search {
text_search.start_worker();
}
thumbnails.start_similarity_warmup();
let initial_thumbnail_manager = Arc::clone(&thumbnails);
tokio::spawn(async move {
initial_thumbnail_manager
@@ -102,6 +143,7 @@ async fn rescan_gallery(state: AppState, config: Config) {
let new_count = updated.images.len();
state.thumbnails.reconcile(&updated.images);
state.thumbnails.start_similarity_warmup();
*state.index.write().await = Arc::new(updated);
state.thumbnails.cleanup_stale().await;
info!(old_count, new_count, "gallery changes indexed");
src/semantic.rs
+1040
View File
File diff suppressed because it is too large. Load diff
src/similarity.rs
+811
View File
@@ -0,0 +1,811 @@
//! Compact, local image descriptors used by the "find similar" endpoint.
//!
//! The descriptor deliberately stays CPU-only and model-free so it remains a
//! dependable fallback when no semantic vision model is configured. It mixes
//! brightness-invariant structure, a perceptual hash, colour distribution,
//! spatial colour and edge orientation. Descriptor sidecars are persisted
//! beside thumbnails so an application restart does not require decoding the
//! whole gallery again.
use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::Write,
sync::atomic::{AtomicU64, Ordering},
sync::{Mutex, OnceLock},
};
use anyhow::{Context, Result, anyhow, bail};
use image::{DynamicImage, ImageReader, imageops::FilterType};
#[cfg(test)]
use image::{ImageDecoder, metadata::Orientation};
use tracing::warn;
use crate::gallery::ImageRecord;
const SAMPLE_EDGE: u32 = 32;
const PIXEL_COUNT: usize = (SAMPLE_EDGE * SAMPLE_EDGE) as usize;
const COLOUR_HUE_BINS: usize = 12;
const COLOUR_SATURATION_BINS: usize = 3;
const COLOUR_VALUE_BINS: usize = 3;
const COLOUR_FEATURES: usize = COLOUR_HUE_BINS * COLOUR_SATURATION_BINS * COLOUR_VALUE_BINS;
const SPATIAL_EDGE: usize = 4;
const SPATIAL_FEATURES: usize = SPATIAL_EDGE * SPATIAL_EDGE * 3;
const EDGE_ORIENTATION_BINS: usize = 8;
const EDGE_FEATURES: usize = SPATIAL_EDGE * SPATIAL_EDGE * EDGE_ORIENTATION_BINS;
const HASH_EDGE: usize = 8;
const HASH_BITS: f32 = (HASH_EDGE * HASH_EDGE - 1) as f32;
const CACHE_LIMIT: usize = 16_384;
const SIGNATURE_MAGIC: &[u8; 8] = b"PXHFSIM2";
const SIGNATURE_EXTENSION: &str = "sim2";
const SIGNATURE_LENGTH: u64 = (SIGNATURE_MAGIC.len()
+ PIXEL_COUNT
+ COLOUR_FEATURES * 2
+ SPATIAL_FEATURES
+ EDGE_FEATURES * 2
+ 8
+ 3
+ 4 * 3) as u64;
/// Scores below this value are generally a coincidental palette/layout match,
/// rather than a useful recommendation. Keep this in one place so a future
/// semantic model can calibrate its score before entering the same pipeline.
pub(crate) const MIN_SIMILARITY_SCORE: f32 = 0.58;
#[derive(Clone)]
pub(crate) struct ImageSignature {
/// Z-normalised luminance samples. This makes structure largely invariant
/// to global exposure and contrast changes.
normalized_luminance: [i8; PIXEL_COUNT],
/// Joint HSV histogram; sample counts always add up to `PIXEL_COUNT`.
colour_histogram: [u16; COLOUR_FEATURES],
/// Mean Y/Cb/Cr values for a 4x4 spatial grid.
spatial_colour: [u8; SPATIAL_FEATURES],
/// Per-cell Sobel orientation histograms, each normalised to 0..65535.
edge_histogram: [u16; EDGE_FEATURES],
perceptual_hash: u64,
average_colour: [u8; 3],
contrast: f32,
edge_energy: f32,
aspect_ratio: f32,
}
#[derive(Clone, Copy)]
pub(crate) struct DiversityFingerprint {
perceptual_hash: u64,
average_colour: [u8; 3],
aspect_ratio: f32,
}
impl ImageSignature {
pub(crate) fn diversity_fingerprint(&self) -> DiversityFingerprint {
DiversityFingerprint {
perceptual_hash: self.perceptual_hash,
average_colour: self.average_colour,
aspect_ratio: self.aspect_ratio,
}
}
}
static SIGNATURE_CACHE: OnceLock<Mutex<HashMap<String, ImageSignature>>> = OnceLock::new();
static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
static DCT_BASIS: OnceLock<[[f32; SAMPLE_EDGE as usize]; HASH_EDGE]> = OnceLock::new();
fn signature_cache() -> &'static Mutex<HashMap<String, ImageSignature>> {
SIGNATURE_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn lock_cache() -> std::sync::MutexGuard<'static, HashMap<String, ImageSignature>> {
signature_cache()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn cache_signature(id: &str, signature: ImageSignature) {
let mut cache = lock_cache();
if cache.len() >= CACHE_LIMIT {
// IDs include the file fingerprint, so evicting any old entry is safe.
if let Some(evicted) = cache.keys().next().cloned() {
cache.remove(&evicted);
}
}
cache.insert(id.to_owned(), signature);
}
/// Return a cached signature, decoding the original only in focused unit tests.
#[cfg(test)]
pub(crate) fn signature_for(record: &ImageRecord) -> Result<ImageSignature> {
record.ensure_source_is_current()?;
let reader = ImageReader::open(&record.path)
.with_context(|| format!("cannot open image {}", record.path.display()))?
.with_guessed_format()
.with_context(|| format!("cannot identify image {}", record.path.display()))?;
let mut decoder = reader
.into_decoder()
.with_context(|| format!("cannot decode image {}", record.path.display()))?;
let orientation = decoder.orientation().unwrap_or(Orientation::NoTransforms);
let mut image = DynamicImage::from_decoder(decoder)
.with_context(|| format!("cannot decode image {}", record.path.display()))?;
image.apply_orientation(orientation);
let signature = signature_from_image(record, &image);
record.ensure_source_is_current()?;
Ok(signature)
}
/// Load a persisted descriptor or derive one from the already-oriented,
/// bounded thumbnail. A broken sidecar is simply regenerated.
pub(crate) fn signature_for_thumbnail(
record: &ImageRecord,
thumbnail: &std::path::Path,
) -> Result<ImageSignature> {
if let Some(signature) = lock_cache().get(&record.id).cloned() {
return Ok(signature);
}
let sidecar = signature_path(thumbnail);
if let Ok(signature) = read_signature(&sidecar) {
cache_signature(&record.id, signature.clone());
return Ok(signature);
}
let image = ImageReader::open(thumbnail)
.with_context(|| format!("cannot open thumbnail {}", thumbnail.display()))?
.with_guessed_format()
.with_context(|| format!("cannot identify thumbnail {}", thumbnail.display()))?
.decode()
.with_context(|| format!("cannot decode thumbnail {}", thumbnail.display()))?;
let signature = signature_from_image(record, &image);
if let Err(error) = write_signature(&sidecar, &signature) {
warn!(path = %sidecar.display(), %error, "cannot persist image similarity descriptor");
}
cache_signature(&record.id, signature.clone());
Ok(signature)
}
/// Persist the descriptor while a thumbnail is already decoded. This makes
/// similarity search cheap even before the first request after a fresh scan.
pub(crate) fn cache_signature_for_thumbnail(
record: &ImageRecord,
image: &DynamicImage,
thumbnail: &std::path::Path,
) -> Result<()> {
let signature = signature_from_image(record, image);
write_signature(&signature_path(thumbnail), &signature)?;
cache_signature(&record.id, signature);
Ok(())
}
pub(crate) fn remove_signature_for_thumbnail(thumbnail: &std::path::Path) -> std::io::Result<()> {
match fs::remove_file(signature_path(thumbnail)) {
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
result => result,
}
}
pub(crate) fn is_signature_sidecar(path: &std::path::Path) -> bool {
path.extension().and_then(|extension| extension.to_str()) == Some(SIGNATURE_EXTENSION)
}
pub(crate) fn has_signature_sidecar(thumbnail: &std::path::Path) -> bool {
fs::metadata(signature_path(thumbnail))
.is_ok_and(|metadata| metadata.is_file() && metadata.len() == SIGNATURE_LENGTH)
}
fn signature_from_image(record: &ImageRecord, image: &DynamicImage) -> ImageSignature {
let sampled = image
.resize_exact(SAMPLE_EDGE, SAMPLE_EDGE, FilterType::Triangle)
.to_rgb8();
let mut luminance = [0_u8; PIXEL_COUNT];
let mut colour_histogram = [0_u16; COLOUR_FEATURES];
let mut spatial_sums = [[0_u32; 3]; SPATIAL_EDGE * SPATIAL_EDGE];
let mut spatial_counts = [0_u32; SPATIAL_EDGE * SPATIAL_EDGE];
let mut average_sums = [0_u32; 3];
for (x, y, pixel) in sampled.enumerate_pixels() {
let [red, green, blue] = pixel.0;
let index = (y * SAMPLE_EDGE + x) as usize;
luminance[index] = rgb_luminance(red, green, blue);
let (hue, saturation, value) = rgb_to_hsv(red, green, blue);
let hue_bin = ((hue * COLOUR_HUE_BINS as f32) as usize).min(COLOUR_HUE_BINS - 1);
let saturation_bin =
((saturation * COLOUR_SATURATION_BINS as f32) as usize).min(COLOUR_SATURATION_BINS - 1);
let value_bin = ((value * COLOUR_VALUE_BINS as f32) as usize).min(COLOUR_VALUE_BINS - 1);
let colour_index =
(hue_bin * COLOUR_SATURATION_BINS + saturation_bin) * COLOUR_VALUE_BINS + value_bin;
colour_histogram[colour_index] += 1;
let cell_x = (x as usize * SPATIAL_EDGE / SAMPLE_EDGE as usize).min(SPATIAL_EDGE - 1);
let cell_y = (y as usize * SPATIAL_EDGE / SAMPLE_EDGE as usize).min(SPATIAL_EDGE - 1);
let cell = cell_y * SPATIAL_EDGE + cell_x;
let [y_colour, cb, cr] = rgb_to_ycbcr(red, green, blue);
spatial_sums[cell][0] += u32::from(y_colour);
spatial_sums[cell][1] += u32::from(cb);
spatial_sums[cell][2] += u32::from(cr);
spatial_counts[cell] += 1;
average_sums[0] += u32::from(y_colour);
average_sums[1] += u32::from(cb);
average_sums[2] += u32::from(cr);
}
let mean = luminance.iter().map(|value| f32::from(*value)).sum::<f32>() / PIXEL_COUNT as f32;
let variance = luminance
.iter()
.map(|value| {
let difference = f32::from(*value) - mean;
difference * difference
})
.sum::<f32>()
/ PIXEL_COUNT as f32;
let standard_deviation = variance.sqrt();
let normalization_scale = standard_deviation.max(8.0);
let normalized_luminance = std::array::from_fn(|index| {
(((f32::from(luminance[index]) - mean) / normalization_scale) * 36.0)
.round()
.clamp(-127.0, 127.0) as i8
});
let mut spatial_colour = [0_u8; SPATIAL_FEATURES];
for cell in 0..spatial_counts.len() {
let count = spatial_counts[cell].max(1);
for channel in 0..3 {
spatial_colour[cell * 3 + channel] =
(spatial_sums[cell][channel] / count).min(255) as u8;
}
}
let (edge_histogram, edge_energy) = edge_descriptor(&luminance);
let average_colour =
std::array::from_fn(|channel| (average_sums[channel] / PIXEL_COUNT as u32).min(255) as u8);
ImageSignature {
normalized_luminance,
colour_histogram,
spatial_colour,
edge_histogram,
perceptual_hash: perceptual_hash(&luminance),
average_colour,
contrast: standard_deviation / 255.0,
edge_energy,
aspect_ratio: record.width as f32 / record.height.max(1) as f32,
}
}
fn rgb_luminance(red: u8, green: u8, blue: u8) -> u8 {
(f32::from(red) * 0.299 + f32::from(green) * 0.587 + f32::from(blue) * 0.114).round() as u8
}
fn rgb_to_ycbcr(red: u8, green: u8, blue: u8) -> [u8; 3] {
let red = f32::from(red);
let green = f32::from(green);
let blue = f32::from(blue);
let y = red * 0.299 + green * 0.587 + blue * 0.114;
let cb = 128.0 - red * 0.168_736 - green * 0.331_264 + blue * 0.5;
let cr = 128.0 + red * 0.5 - green * 0.418_688 - blue * 0.081_312;
[y, cb, cr].map(|value| value.round().clamp(0.0, 255.0) as u8)
}
fn rgb_to_hsv(red: u8, green: u8, blue: u8) -> (f32, f32, f32) {
let red = f32::from(red) / 255.0;
let green = f32::from(green) / 255.0;
let blue = f32::from(blue) / 255.0;
let maximum = red.max(green).max(blue);
let minimum = red.min(green).min(blue);
let delta = maximum - minimum;
let saturation = if maximum <= f32::EPSILON {
0.0
} else {
delta / maximum
};
let hue = if delta <= f32::EPSILON {
0.0
} else if maximum == red {
((green - blue) / delta).rem_euclid(6.0) / 6.0
} else if maximum == green {
(((blue - red) / delta) + 2.0) / 6.0
} else {
(((red - green) / delta) + 4.0) / 6.0
};
(hue, saturation, maximum)
}
fn edge_descriptor(luminance: &[u8; PIXEL_COUNT]) -> ([u16; EDGE_FEATURES], f32) {
let mut histogram = [0.0_f32; EDGE_FEATURES];
let mut energy = 0.0_f32;
for y in 1..SAMPLE_EDGE as usize - 1 {
for x in 1..SAMPLE_EDGE as usize - 1 {
let sample = |offset_x: isize, offset_y: isize| -> f32 {
let x = (x as isize + offset_x) as usize;
let y = (y as isize + offset_y) as usize;
f32::from(luminance[y * SAMPLE_EDGE as usize + x])
};
let gradient_x = -sample(-1, -1) + sample(1, -1) - 2.0 * sample(-1, 0)
+ 2.0 * sample(1, 0)
- sample(-1, 1)
+ sample(1, 1);
let gradient_y = -sample(-1, -1) - 2.0 * sample(0, -1) - sample(1, -1)
+ sample(-1, 1)
+ 2.0 * sample(0, 1)
+ sample(1, 1);
let magnitude = gradient_x.hypot(gradient_y);
energy += magnitude;
if magnitude < 4.0 {
continue;
}
let angle = gradient_y
.atan2(gradient_x)
.rem_euclid(std::f32::consts::PI);
let orientation = ((angle / std::f32::consts::PI * EDGE_ORIENTATION_BINS as f32)
as usize)
.min(EDGE_ORIENTATION_BINS - 1);
let cell_x = (x * SPATIAL_EDGE / SAMPLE_EDGE as usize).min(SPATIAL_EDGE - 1);
let cell_y = (y * SPATIAL_EDGE / SAMPLE_EDGE as usize).min(SPATIAL_EDGE - 1);
let cell = cell_y * SPATIAL_EDGE + cell_x;
histogram[cell * EDGE_ORIENTATION_BINS + orientation] += magnitude;
}
}
let mut normalized = [0_u16; EDGE_FEATURES];
for cell in 0..SPATIAL_EDGE * SPATIAL_EDGE {
let start = cell * EDGE_ORIENTATION_BINS;
let total = histogram[start..start + EDGE_ORIENTATION_BINS]
.iter()
.sum::<f32>();
if total <= f32::EPSILON {
continue;
}
for orientation in 0..EDGE_ORIENTATION_BINS {
normalized[start + orientation] = (histogram[start + orientation] / total
* f32::from(u16::MAX))
.round()
.clamp(0.0, f32::from(u16::MAX)) as u16;
}
}
let interior_pixels = ((SAMPLE_EDGE - 2) * (SAMPLE_EDGE - 2)) as f32;
// Sobel magnitude can reach roughly 4 * 255. Normalising keeps the value
// architecture-independent and easy to persist.
(
normalized,
(energy / interior_pixels / 1020.0).clamp(0.0, 1.0),
)
}
fn perceptual_hash(luminance: &[u8; PIXEL_COUNT]) -> u64 {
let basis = DCT_BASIS.get_or_init(|| {
std::array::from_fn(|frequency| {
std::array::from_fn(|position| {
((std::f32::consts::PI / SAMPLE_EDGE as f32)
* (position as f32 + 0.5)
* frequency as f32)
.cos()
})
})
});
// DCT is separable. Computing the horizontal low frequencies once cuts
// descriptor construction from roughly 65k to 10k multiply-adds.
let mut horizontal = [0.0_f32; SAMPLE_EDGE as usize * HASH_EDGE];
for y in 0..SAMPLE_EDGE as usize {
for horizontal_frequency in 0..HASH_EDGE {
horizontal[y * HASH_EDGE + horizontal_frequency] = (0..SAMPLE_EDGE as usize)
.map(|x| {
f32::from(luminance[y * SAMPLE_EDGE as usize + x])
* basis[horizontal_frequency][x]
})
.sum();
}
}
let mut coefficients = [0.0_f32; HASH_EDGE * HASH_EDGE];
for vertical_frequency in 0..HASH_EDGE {
for horizontal_frequency in 0..HASH_EDGE {
coefficients[vertical_frequency * HASH_EDGE + horizontal_frequency] = (0..SAMPLE_EDGE
as usize)
.map(|y| {
horizontal[y * HASH_EDGE + horizontal_frequency] * basis[vertical_frequency][y]
})
.sum();
}
}
let mut frequencies = coefficients[1..].to_vec();
frequencies.sort_unstable_by(f32::total_cmp);
let median = frequencies[frequencies.len() / 2];
coefficients[1..]
.iter()
.enumerate()
.fold(0_u64, |hash, (index, coefficient)| {
hash | (u64::from(*coefficient > median) << index)
})
}
/// Score two descriptors from 0 (unrelated) to 1 (near-identical).
pub(crate) fn similarity_score(first: &ImageSignature, second: &ImageSignature) -> f32 {
let colour = colour_similarity(first, second);
let spatial = spatial_colour_similarity(first, second);
let structure = structure_similarity(first, second);
let hash = hash_similarity(first.perceptual_hash, second.perceptual_hash);
let edges = edge_similarity(first, second);
let aspect = aspect_similarity(first.aspect_ratio, second.aspect_ratio);
// Flat graphics and scans have little structural evidence, so dynamically
// move their weight to colour instead of letting an all-zero edge map look
// like a perfect match.
let contrast_reliability = (first.contrast.min(second.contrast) / 0.18).clamp(0.0, 1.0);
let edge_reliability = (first.edge_energy.min(second.edge_energy) / 0.08).clamp(0.0, 1.0);
let colour_weight = 0.32;
let spatial_weight = 0.22;
let structure_weight = 0.20 * contrast_reliability;
let hash_weight = 0.16 * contrast_reliability;
let edge_weight = 0.18 * edge_reliability;
let aspect_weight = 0.04;
let weight = colour_weight
+ spatial_weight
+ structure_weight
+ hash_weight
+ edge_weight
+ aspect_weight;
((colour * colour_weight
+ spatial * spatial_weight
+ structure * structure_weight
+ hash * hash_weight
+ edges * edge_weight
+ aspect * aspect_weight)
/ weight.max(f32::EPSILON))
.clamp(0.0, 1.0)
}
fn colour_similarity(first: &ImageSignature, second: &ImageSignature) -> f32 {
first
.colour_histogram
.iter()
.zip(&second.colour_histogram)
.map(|(left, right)| f32::from((*left).min(*right)))
.sum::<f32>()
/ PIXEL_COUNT as f32
}
fn spatial_colour_similarity(first: &ImageSignature, second: &ImageSignature) -> f32 {
let distance = first
.spatial_colour
.chunks_exact(3)
.zip(second.spatial_colour.chunks_exact(3))
.map(|(left, right)| {
let luminance = f32::from(left[0].abs_diff(right[0])) / 255.0;
let chroma = (f32::from(left[1].abs_diff(right[1]))
+ f32::from(left[2].abs_diff(right[2])))
/ (2.0 * 255.0);
luminance * 0.45 + chroma * 0.55
})
.sum::<f32>()
/ (SPATIAL_EDGE * SPATIAL_EDGE) as f32;
(-4.0 * distance).exp()
}
fn structure_similarity(first: &ImageSignature, second: &ImageSignature) -> f32 {
let mut dot = 0.0_f32;
let mut first_energy = 0.0_f32;
let mut second_energy = 0.0_f32;
for (left, right) in first
.normalized_luminance
.iter()
.zip(&second.normalized_luminance)
{
let left = f32::from(*left);
let right = f32::from(*right);
dot += left * right;
first_energy += left * left;
second_energy += right * right;
}
if first_energy <= f32::EPSILON || second_energy <= f32::EPSILON {
return 0.5;
}
let correlation = dot / (first_energy * second_energy).sqrt();
((correlation.clamp(-1.0, 1.0) + 1.0) * 0.5).clamp(0.0, 1.0)
}
fn hash_similarity(first: u64, second: u64) -> f32 {
(1.0 - (first ^ second).count_ones() as f32 / HASH_BITS).clamp(0.0, 1.0)
}
fn edge_similarity(first: &ImageSignature, second: &ImageSignature) -> f32 {
let mut similarity = 0.0_f32;
let mut compared_cells = 0usize;
for cell in 0..SPATIAL_EDGE * SPATIAL_EDGE {
let start = cell * EDGE_ORIENTATION_BINS;
let first_cell = &first.edge_histogram[start..start + EDGE_ORIENTATION_BINS];
let second_cell = &second.edge_histogram[start..start + EDGE_ORIENTATION_BINS];
let first_total = first_cell
.iter()
.map(|value| u32::from(*value))
.sum::<u32>();
let second_total = second_cell
.iter()
.map(|value| u32::from(*value))
.sum::<u32>();
if first_total == 0 && second_total == 0 {
continue;
}
compared_cells += 1;
similarity += first_cell
.iter()
.zip(second_cell)
.map(|(left, right)| f32::from((*left).min(*right)))
.sum::<f32>()
/ f32::from(u16::MAX);
}
if compared_cells == 0 {
0.5
} else {
(similarity / compared_cells as f32).clamp(0.0, 1.0)
}
}
fn aspect_similarity(first: f32, second: f32) -> f32 {
(-(first / second.max(0.001)).ln().abs()).exp()
}
/// Cheap near-duplicate estimate used only to diversify the first page. It is
/// intentionally stricter than the main relevance score.
pub(crate) fn redundancy_score(first: DiversityFingerprint, second: DiversityFingerprint) -> f32 {
let hash = hash_similarity(first.perceptual_hash, second.perceptual_hash);
let colour_distance = first
.average_colour
.iter()
.zip(second.average_colour)
.map(|(left, right)| f32::from(left.abs_diff(right)) / 255.0)
.sum::<f32>()
/ 3.0;
let colour = (-5.0 * colour_distance).exp();
let aspect = aspect_similarity(first.aspect_ratio, second.aspect_ratio);
(hash * 0.78 + colour * 0.17 + aspect * 0.05).clamp(0.0, 1.0)
}
fn signature_path(thumbnail: &std::path::Path) -> std::path::PathBuf {
thumbnail.with_extension(SIGNATURE_EXTENSION)
}
fn write_signature(path: &std::path::Path, signature: &ImageSignature) -> Result<()> {
let parent = path
.parent()
.with_context(|| format!("invalid descriptor cache path: {}", path.display()))?;
fs::create_dir_all(parent)?;
let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
let temporary = parent.join(format!(
".{}.{}.{}.tmp",
path.file_stem()
.and_then(|value| value.to_str())
.unwrap_or("image"),
std::process::id(),
sequence
));
let result = (|| -> Result<()> {
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(&temporary)?;
file.write_all(SIGNATURE_MAGIC)?;
file.write_all(&signature.normalized_luminance.map(|value| value as u8))?;
for value in signature.colour_histogram {
file.write_all(&value.to_le_bytes())?;
}
file.write_all(&signature.spatial_colour)?;
for value in signature.edge_histogram {
file.write_all(&value.to_le_bytes())?;
}
file.write_all(&signature.perceptual_hash.to_le_bytes())?;
file.write_all(&signature.average_colour)?;
file.write_all(&signature.contrast.to_le_bytes())?;
file.write_all(&signature.edge_energy.to_le_bytes())?;
file.write_all(&signature.aspect_ratio.to_le_bytes())?;
drop(file);
fs::rename(&temporary, path)?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
}
result
}
fn read_signature(path: &std::path::Path) -> Result<ImageSignature> {
let bytes = fs::read(path)?;
let mut cursor = 0usize;
let magic = take_bytes::<8>(&bytes, &mut cursor)?;
if &magic != SIGNATURE_MAGIC {
bail!("unknown descriptor cache version");
}
let normalized_luminance =
take_bytes::<PIXEL_COUNT>(&bytes, &mut cursor)?.map(|value| value as i8);
let mut colour_histogram = [0_u16; COLOUR_FEATURES];
for value in &mut colour_histogram {
*value = u16::from_le_bytes(take_bytes::<2>(&bytes, &mut cursor)?);
}
let spatial_colour = take_bytes::<SPATIAL_FEATURES>(&bytes, &mut cursor)?;
let mut edge_histogram = [0_u16; EDGE_FEATURES];
for value in &mut edge_histogram {
*value = u16::from_le_bytes(take_bytes::<2>(&bytes, &mut cursor)?);
}
let perceptual_hash = u64::from_le_bytes(take_bytes::<8>(&bytes, &mut cursor)?);
let average_colour = take_bytes::<3>(&bytes, &mut cursor)?;
let contrast = f32::from_le_bytes(take_bytes::<4>(&bytes, &mut cursor)?);
let edge_energy = f32::from_le_bytes(take_bytes::<4>(&bytes, &mut cursor)?);
let aspect_ratio = f32::from_le_bytes(take_bytes::<4>(&bytes, &mut cursor)?);
if cursor != bytes.len()
|| !contrast.is_finite()
|| !edge_energy.is_finite()
|| !aspect_ratio.is_finite()
|| aspect_ratio <= 0.0
{
bail!("invalid descriptor cache contents");
}
Ok(ImageSignature {
normalized_luminance,
colour_histogram,
spatial_colour,
edge_histogram,
perceptual_hash,
average_colour,
contrast,
edge_energy,
aspect_ratio,
})
}
fn take_bytes<const LENGTH: usize>(bytes: &[u8], cursor: &mut usize) -> Result<[u8; LENGTH]> {
let end = cursor
.checked_add(LENGTH)
.ok_or_else(|| anyhow!("descriptor cache length overflow"))?;
let value = bytes
.get(*cursor..end)
.ok_or_else(|| anyhow!("truncated descriptor cache"))?
.try_into()
.map_err(|_| anyhow!("invalid descriptor cache field"))?;
*cursor = end;
Ok(value)
}
#[cfg(test)]
mod tests {
use super::*;
use std::{fs, path::PathBuf, time::UNIX_EPOCH};
use crate::gallery::ImageRecord;
fn record(path: PathBuf, id: &str, width: u32, height: u32) -> ImageRecord {
let metadata = fs::metadata(&path).unwrap();
let modified_ns = metadata
.modified()
.unwrap()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
ImageRecord {
id: id.to_owned(),
path,
relative_path: id.to_owned(),
search_key: id.to_owned(),
name: id.to_owned(),
album: String::new(),
width,
height,
size: metadata.len(),
modified_ms: 0,
modified_ns,
}
}
#[test]
fn identical_images_score_higher_than_different_images() {
let temp = tempfile::tempdir().unwrap();
let first_path = temp.path().join("first.png");
let same_path = temp.path().join("same.png");
let different_path = temp.path().join("different.png");
image::RgbImage::from_pixel(48, 32, image::Rgb([220, 80, 40]))
.save(&first_path)
.unwrap();
fs::copy(&first_path, &same_path).unwrap();
image::RgbImage::from_pixel(48, 32, image::Rgb([20, 40, 220]))
.save(&different_path)
.unwrap();
let first = record(first_path, "colour-first", 48, 32);
let same = record(same_path, "colour-same", 48, 32);
let different = record(different_path, "colour-different", 48, 32);
let first_signature = signature_for(&first).unwrap();
let same_signature = signature_for(&same).unwrap();
let different_signature = signature_for(&different).unwrap();
assert!(
similarity_score(&first_signature, &same_signature)
> similarity_score(&first_signature, &different_signature)
);
assert!(similarity_score(&first_signature, &same_signature) > 0.99);
assert!(similarity_score(&first_signature, &different_signature) < MIN_SIMILARITY_SCORE);
}
#[test]
fn descriptor_sidecars_round_trip_and_reject_corruption() {
let temp = tempfile::tempdir().unwrap();
let image_path = temp.path().join("source.png");
let sidecar = temp.path().join("source.sim2");
image::RgbImage::from_fn(64, 48, |x, y| {
image::Rgb([(x * 3) as u8, (y * 5) as u8, ((x + y) * 2) as u8])
})
.save(&image_path)
.unwrap();
let source = record(image_path, "sidecar-source", 64, 48);
let signature = signature_for(&source).unwrap();
write_signature(&sidecar, &signature).unwrap();
let restored = read_signature(&sidecar).unwrap();
assert_eq!(similarity_score(&signature, &restored), 1.0);
fs::write(&sidecar, b"not a signature").unwrap();
assert!(read_signature(&sidecar).is_err());
}
#[test]
fn exposure_and_small_crops_keep_the_same_scene_ahead() {
let temp = tempfile::tempdir().unwrap();
let source_path = temp.path().join("scene-source.png");
let exposure_path = temp.path().join("scene-exposure.png");
let crop_path = temp.path().join("scene-crop.png");
let unrelated_path = temp.path().join("scene-unrelated.png");
let source_image = image::RgbImage::from_fn(96, 72, |x, y| {
let circle = (x as i32 - 30).pow(2) + (y as i32 - 36).pow(2) < 18_i32.pow(2);
if circle {
image::Rgb([218, 82, 43])
} else if x > 56 && (16..60).contains(&y) {
image::Rgb([42, 166, 112])
} else {
image::Rgb([
28_u8.saturating_add((x / 2) as u8),
42_u8.saturating_add((y / 2) as u8),
92_u8.saturating_add(((x + y) / 5) as u8),
])
}
});
let exposure_image =
image::RgbImage::from_fn(96, 72, |x, y| {
image::Rgb(source_image.get_pixel(x, y).0.map(|channel| {
(f32::from(channel) * 0.82 + 24.0).round().clamp(0.0, 255.0) as u8
}))
});
let crop_image = image::imageops::crop_imm(&source_image, 4, 3, 88, 66).to_image();
let unrelated_image = image::RgbImage::from_fn(96, 72, |x, y| {
let stripe = ((x / 12) + (y / 9)) % 2 == 0;
if stripe {
image::Rgb([218, 82, 43])
} else {
image::Rgb([42, 166, 112])
}
});
source_image.save(&source_path).unwrap();
exposure_image.save(&exposure_path).unwrap();
crop_image.save(&crop_path).unwrap();
unrelated_image.save(&unrelated_path).unwrap();
let source = record(source_path, "robust-source", 96, 72);
let exposure = record(exposure_path, "robust-exposure", 96, 72);
let crop = record(crop_path, "robust-crop", 88, 66);
let unrelated = record(unrelated_path, "robust-unrelated", 96, 72);
let source_signature = signature_for(&source).unwrap();
let exposure_score =
similarity_score(&source_signature, &signature_for(&exposure).unwrap());
let crop_score = similarity_score(&source_signature, &signature_for(&crop).unwrap());
let unrelated_score =
similarity_score(&source_signature, &signature_for(&unrelated).unwrap());
assert!(
exposure_score >= MIN_SIMILARITY_SCORE,
"exposure score was {exposure_score}"
);
assert!(
crop_score >= MIN_SIMILARITY_SCORE,
"crop score was {crop_score}"
);
assert!(
exposure_score > unrelated_score && crop_score > unrelated_score,
"exposure={exposure_score}, crop={crop_score}, unrelated={unrelated_score}"
);
}
}
src/text_search.rs
+885
View File
@@ -0,0 +1,885 @@
//! Optional Chinese-CLIP embeddings for natural-language image search.
//!
//! The model, vocabulary, image embeddings, and text queries stay local. A
//! single CPU worker builds compact quantised image vectors beside thumbnails;
//! text queries use the same shared embedding space and are cached in memory.
use std::{
collections::{HashMap, HashSet, VecDeque},
fmt::Write as _,
fs::{self, File, OpenOptions},
io::{Read, Write},
path::{Path, PathBuf},
sync::{
Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
atomic::{AtomicBool, AtomicU64, Ordering},
},
};
use anyhow::{Context, Result, anyhow, bail};
use candle_core::{DType, Device, IndexOp, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::chinese_clip::{ChineseClipConfig, ChineseClipModel};
use image::{DynamicImage, GenericImageView, ImageReader, imageops::FilterType};
use serde::Serialize;
use tokenizers::{
Tokenizer, TruncationParams, models::wordpiece::WordPiece, normalizers::bert::BertNormalizer,
pre_tokenizers::bert::BertPreTokenizer, processors::bert::BertProcessing,
};
use tokio::sync::Notify;
use tracing::{info, warn};
const MODEL_NAME: &str = "Chinese-CLIP ViT-B/16";
const MODEL_EDGE: u32 = 224;
const TEXT_CONTEXT_LENGTH: usize = 52;
const EMBEDDING_SIZE: usize = 512;
const EMBEDDING_MAGIC: &[u8; 8] = b"PXHFCN1\0";
const EMBEDDING_EXTENSION: &str = "cnclip";
const MODEL_FINGERPRINT_BYTES: usize = 16;
const EMBEDDING_CHECKSUM_BYTES: usize = 16;
const EMBEDDING_LENGTH: u64 =
(EMBEDDING_MAGIC.len() + MODEL_FINGERPRINT_BYTES + EMBEDDING_SIZE + EMBEDDING_CHECKSUM_BYTES)
as u64;
const QUERY_CACHE_LIMIT: usize = 16;
static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
#[derive(Clone, Debug)]
pub(crate) struct TextSearchEmbedding {
values: [i8; EMBEDDING_SIZE],
norm: f32,
}
impl TextSearchEmbedding {
fn quantize(values: &[f32]) -> Result<Self> {
if values.len() != EMBEDDING_SIZE || values.iter().any(|value| !value.is_finite()) {
bail!("text-search model returned an invalid embedding");
}
let norm = values.iter().map(|value| value * value).sum::<f32>().sqrt();
if norm <= f32::EPSILON {
bail!("text-search model returned an empty embedding");
}
let values = std::array::from_fn(|index| {
(values[index] / norm * 127.0).round().clamp(-127.0, 127.0) as i8
});
Self::from_quantized(values)
}
fn from_quantized(values: [i8; EMBEDDING_SIZE]) -> Result<Self> {
let norm = values
.iter()
.map(|value| f32::from(*value).powi(2))
.sum::<f32>()
.sqrt();
if norm <= f32::EPSILON {
bail!("text-search embedding is empty");
}
Ok(Self { values, norm })
}
#[cfg(test)]
pub(crate) fn for_test(values: impl IntoIterator<Item = (usize, i8)>) -> Self {
let mut embedding = [0_i8; EMBEDDING_SIZE];
for (index, value) in values {
embedding[index] = value;
}
Self::from_quantized(embedding).expect("test embedding must not be empty")
}
}
pub(crate) fn text_image_similarity(
query: &TextSearchEmbedding,
image: &TextSearchEmbedding,
) -> f32 {
let dot = query
.values
.iter()
.zip(&image.values)
.map(|(left, right)| f32::from(*left) * f32::from(*right))
.sum::<f32>();
(dot / (query.norm * image.norm)).clamp(-1.0, 1.0)
}
pub(crate) trait TextImageModel: Send {
fn embed_thumbnail(&mut self, thumbnail: &Path) -> Result<TextSearchEmbedding>;
fn embed_text(&mut self, text: &str) -> Result<TextSearchEmbedding>;
}
pub(crate) struct TextSearchIndex {
model: Arc<Mutex<Box<dyn TextImageModel>>>,
model_fingerprint: [u8; MODEL_FINGERPRINT_BYTES],
model_token: String,
entries: RwLock<HashMap<String, Arc<TextSearchEntry>>>,
queue: TextSearchQueue,
query_cache: Mutex<VecDeque<(String, Arc<TextSearchEmbedding>)>>,
worker_started: AtomicBool,
revision: AtomicU64,
}
struct TextSearchEntry {
thumbnail: PathBuf,
state: Mutex<TextSearchState>,
notify: Notify,
}
enum TextSearchState {
Pending,
Processing,
Ready(Arc<TextSearchEmbedding>),
Failed,
Removed,
}
#[derive(Default)]
struct TextSearchQueue {
state: Mutex<TextSearchQueueState>,
notify: Notify,
}
#[derive(Default)]
struct TextSearchQueueState {
urgent: VecDeque<String>,
background: VecDeque<String>,
queued: HashSet<String>,
}
#[derive(Clone, Copy, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct TextSearchStatus {
pub enabled: bool,
pub total: usize,
pub ready: usize,
pub queued: usize,
pub processing: usize,
pub failed: usize,
pub background_complete: bool,
}
impl TextSearchStatus {
pub(crate) fn disabled() -> Self {
Self {
enabled: false,
total: 0,
ready: 0,
queued: 0,
processing: 0,
failed: 0,
background_complete: true,
}
}
}
fn mutex_lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn read_lock<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
lock.read().unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn write_lock<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
lock.write()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
impl TextSearchIndex {
pub(crate) fn load(model_path: &Path, vocabulary_path: &Path) -> Result<Arc<Self>> {
let fingerprint = fingerprint_files(&[model_path, vocabulary_path])?;
let model = ChineseClip::load(model_path, vocabulary_path).with_context(|| {
format!(
"cannot load natural-language search model: {}",
model_path.display()
)
})?;
let index = Self::with_model(Box::new(model), fingerprint);
info!(
model = MODEL_NAME,
fingerprint = %index.model_token,
"local natural-language image search enabled"
);
Ok(index)
}
#[cfg(test)]
pub(crate) fn with_test_model(
model: Box<dyn TextImageModel>,
model_fingerprint: [u8; MODEL_FINGERPRINT_BYTES],
) -> Arc<Self> {
Self::with_model(model, model_fingerprint)
}
fn with_model(
model: Box<dyn TextImageModel>,
model_fingerprint: [u8; MODEL_FINGERPRINT_BYTES],
) -> Arc<Self> {
Arc::new(Self {
model: Arc::new(Mutex::new(model)),
model_fingerprint,
model_token: fingerprint_token(&model_fingerprint),
entries: RwLock::new(HashMap::new()),
queue: TextSearchQueue::default(),
query_cache: Mutex::new(VecDeque::new()),
worker_started: AtomicBool::new(false),
revision: AtomicU64::new(0),
})
}
pub(crate) fn start_worker(self: &Arc<Self>) {
if self.worker_started.swap(true, Ordering::AcqRel) {
return;
}
let index = Arc::clone(self);
tokio::spawn(async move {
index.worker_loop().await;
});
}
pub(crate) fn reconcile(&self, valid_ids: &HashSet<&str>) {
{
let mut entries = write_lock(&self.entries);
entries.retain(|id, entry| {
let retained = valid_ids.contains(id.as_str());
if !retained {
*mutex_lock(&entry.state) = TextSearchState::Removed;
entry.notify.notify_waiters();
}
retained
});
}
self.queue.retain(valid_ids);
}
pub(crate) fn enqueue(&self, id: &str, thumbnail: &Path, urgent: bool) {
let (entry, inserted) = {
let mut entries = write_lock(&self.entries);
match entries.get(id) {
Some(entry) => (Arc::clone(entry), false),
None => {
let entry = Arc::new(TextSearchEntry {
thumbnail: thumbnail.to_owned(),
state: Mutex::new(TextSearchState::Pending),
notify: Notify::new(),
});
entries.insert(id.to_owned(), Arc::clone(&entry));
(entry, true)
}
}
};
let pending = matches!(*mutex_lock(&entry.state), TextSearchState::Pending);
if !inserted && !pending {
return;
}
if urgent {
self.queue.push_urgent(id.to_owned());
} else {
self.queue.push_background(id.to_owned());
}
}
pub(crate) fn embedding(&self, id: &str) -> Option<Arc<TextSearchEmbedding>> {
let entry = read_lock(&self.entries).get(id).cloned()?;
let state = mutex_lock(&entry.state);
match &*state {
TextSearchState::Ready(embedding) => Some(Arc::clone(embedding)),
_ => None,
}
}
pub(crate) async fn embed_query(&self, query: &str) -> Result<Arc<TextSearchEmbedding>> {
if let Some(embedding) = self.cached_query(query) {
return Ok(embedding);
}
let model = Arc::clone(&self.model);
let owned_query = query.to_owned();
let embedding =
tokio::task::spawn_blocking(move || mutex_lock(&model).embed_text(&owned_query))
.await
.map_err(|error| anyhow!("text-search query worker stopped: {error}"))??;
let embedding = Arc::new(embedding);
self.cache_query(query.to_owned(), Arc::clone(&embedding));
Ok(embedding)
}
fn cached_query(&self, query: &str) -> Option<Arc<TextSearchEmbedding>> {
let mut cache = mutex_lock(&self.query_cache);
let position = cache.iter().position(|(cached, _)| cached == query)?;
let entry = cache.remove(position)?;
let embedding = Arc::clone(&entry.1);
cache.push_back(entry);
Some(embedding)
}
fn cache_query(&self, query: String, embedding: Arc<TextSearchEmbedding>) {
let mut cache = mutex_lock(&self.query_cache);
if let Some(position) = cache.iter().position(|(cached, _)| cached == &query) {
cache.remove(position);
}
if cache.len() >= QUERY_CACHE_LIMIT {
cache.pop_front();
}
cache.push_back((query, embedding));
}
pub(crate) fn cache_token(&self) -> String {
format!(
"{}-{}",
self.model_token,
self.revision.load(Ordering::Acquire)
)
}
pub(crate) fn status(&self) -> TextSearchStatus {
let entries = read_lock(&self.entries);
let mut ready = 0;
let mut queued = 0;
let mut processing = 0;
let mut failed = 0;
for entry in entries.values() {
match &*mutex_lock(&entry.state) {
TextSearchState::Pending => queued += 1,
TextSearchState::Processing => processing += 1,
TextSearchState::Ready(_) => ready += 1,
TextSearchState::Failed => failed += 1,
TextSearchState::Removed => {}
}
}
let total = entries.len();
TextSearchStatus {
enabled: true,
total,
ready,
queued,
processing,
failed,
background_complete: ready + failed == total && processing == 0,
}
}
async fn worker_loop(self: Arc<Self>) {
loop {
let id = self.queue.pop().await;
let entry = read_lock(&self.entries).get(&id).cloned();
let Some(entry) = entry else {
continue;
};
{
let mut state = mutex_lock(&entry.state);
if !matches!(*state, TextSearchState::Pending) {
continue;
}
*state = TextSearchState::Processing;
}
let thumbnail = entry.thumbnail.clone();
let sidecar = embedding_path(&thumbnail);
let model = Arc::clone(&self.model);
let fingerprint = self.model_fingerprint;
let result = tokio::task::spawn_blocking(move || {
if let Ok(embedding) = read_embedding(&sidecar, &fingerprint) {
return Ok(embedding);
}
let embedding = mutex_lock(&model).embed_thumbnail(&thumbnail)?;
if let Err(error) = write_embedding(&sidecar, &fingerprint, &embedding) {
warn!(path = %sidecar.display(), %error, "cannot persist text-search image embedding");
}
Ok(embedding)
})
.await
.map_err(|error| anyhow!("text-search index worker stopped: {error}"))
.and_then(|result| result);
let mut removed = false;
match result {
Ok(embedding) => {
let mut state = mutex_lock(&entry.state);
if matches!(*state, TextSearchState::Processing) {
*state = TextSearchState::Ready(Arc::new(embedding));
self.revision.fetch_add(1, Ordering::AcqRel);
} else {
removed = matches!(*state, TextSearchState::Removed);
}
}
Err(error) => {
let mut state = mutex_lock(&entry.state);
if matches!(*state, TextSearchState::Processing) {
*state = TextSearchState::Failed;
self.revision.fetch_add(1, Ordering::AcqRel);
warn!(image = %id, %error, "text-search image embedding failed");
} else {
removed = matches!(*state, TextSearchState::Removed);
}
}
}
entry.notify.notify_waiters();
if removed && let Err(error) = remove_embedding_for_thumbnail(&entry.thumbnail) {
warn!(path = %entry.thumbnail.display(), %error, "cannot remove stale text-search embedding");
}
}
}
}
impl TextSearchQueue {
fn retain(&self, valid_ids: &HashSet<&str>) {
let mut state = mutex_lock(&self.state);
state.urgent.retain(|id| valid_ids.contains(id.as_str()));
state
.background
.retain(|id| valid_ids.contains(id.as_str()));
state.queued.retain(|id| valid_ids.contains(id.as_str()));
}
fn push_urgent(&self, id: String) {
let should_notify = {
let mut state = mutex_lock(&self.state);
if state.queued.contains(&id) {
if let Some(position) = state.background.iter().position(|queued| queued == &id) {
state.background.remove(position);
state.urgent.push_back(id);
true
} else {
false
}
} else {
state.queued.insert(id.clone());
state.urgent.push_back(id);
true
}
};
if should_notify {
self.notify.notify_one();
}
}
fn push_background(&self, id: String) {
let should_notify = {
let mut state = mutex_lock(&self.state);
if state.queued.insert(id.clone()) {
state.background.push_back(id);
true
} else {
false
}
};
if should_notify {
self.notify.notify_one();
}
}
async fn pop(&self) -> String {
loop {
let notified = self.notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
if let Some(id) = {
let mut state = mutex_lock(&self.state);
let id = state
.urgent
.pop_front()
.or_else(|| state.background.pop_front());
if let Some(id) = &id {
state.queued.remove(id);
}
id
} {
return id;
}
notified.as_mut().await;
}
}
}
pub(crate) fn embedding_path(thumbnail: &Path) -> PathBuf {
thumbnail.with_extension(EMBEDDING_EXTENSION)
}
pub(crate) fn is_text_search_sidecar(path: &Path) -> bool {
path.extension().and_then(|extension| extension.to_str()) == Some(EMBEDDING_EXTENSION)
}
pub(crate) fn remove_embedding_for_thumbnail(thumbnail: &Path) -> std::io::Result<()> {
match fs::remove_file(embedding_path(thumbnail)) {
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
result => result,
}
}
fn read_embedding(
path: &Path,
expected_fingerprint: &[u8; MODEL_FINGERPRINT_BYTES],
) -> Result<TextSearchEmbedding> {
let metadata = fs::metadata(path)?;
if !metadata.is_file() || metadata.len() != EMBEDDING_LENGTH {
bail!("text-search embedding cache has an unexpected size");
}
let mut file = File::open(path)?;
let mut magic = [0_u8; EMBEDDING_MAGIC.len()];
file.read_exact(&mut magic)?;
if &magic != EMBEDDING_MAGIC {
bail!("text-search embedding cache has an unknown version");
}
let mut fingerprint = [0_u8; MODEL_FINGERPRINT_BYTES];
file.read_exact(&mut fingerprint)?;
if &fingerprint != expected_fingerprint {
bail!("text-search embedding belongs to another model");
}
let mut bytes = [0_u8; EMBEDDING_SIZE];
file.read_exact(&mut bytes)?;
let mut checksum = [0_u8; EMBEDDING_CHECKSUM_BYTES];
file.read_exact(&mut checksum)?;
if checksum != embedding_checksum(&fingerprint, &bytes) {
bail!("text-search embedding cache checksum does not match");
}
TextSearchEmbedding::from_quantized(bytes.map(|value| value as i8))
}
fn write_embedding(
path: &Path,
fingerprint: &[u8; MODEL_FINGERPRINT_BYTES],
embedding: &TextSearchEmbedding,
) -> Result<()> {
let parent = path
.parent()
.with_context(|| format!("invalid text-search cache path: {}", path.display()))?;
fs::create_dir_all(parent)?;
let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
let temporary = parent.join(format!(
".{}.{}.{}.tmp",
path.file_stem()
.and_then(|value| value.to_str())
.unwrap_or("embedding"),
std::process::id(),
sequence
));
let result = (|| -> Result<()> {
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(&temporary)?;
file.write_all(EMBEDDING_MAGIC)?;
file.write_all(fingerprint)?;
let values = embedding.values.map(|value| value as u8);
file.write_all(&values)?;
file.write_all(&embedding_checksum(fingerprint, &values))?;
file.sync_all()?;
drop(file);
fs::rename(&temporary, path)?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
}
result
}
fn embedding_checksum(
fingerprint: &[u8; MODEL_FINGERPRINT_BYTES],
values: &[u8; EMBEDDING_SIZE],
) -> [u8; EMBEDDING_CHECKSUM_BYTES] {
let mut hasher = blake3::Hasher::new();
hasher.update(fingerprint);
hasher.update(values);
let mut checksum = [0_u8; EMBEDDING_CHECKSUM_BYTES];
checksum.copy_from_slice(&hasher.finalize().as_bytes()[..EMBEDDING_CHECKSUM_BYTES]);
checksum
}
fn fingerprint_files(paths: &[&Path]) -> Result<[u8; MODEL_FINGERPRINT_BYTES]> {
let mut hasher = blake3::Hasher::new();
let mut buffer = [0_u8; 128 * 1024];
for path in paths {
let mut file = File::open(path)
.with_context(|| format!("cannot open text-search model file: {}", path.display()))?;
loop {
let read = file.read(&mut buffer)?;
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
}
hasher.update(b"\0");
}
let mut fingerprint = [0_u8; MODEL_FINGERPRINT_BYTES];
fingerprint.copy_from_slice(&hasher.finalize().as_bytes()[..MODEL_FINGERPRINT_BYTES]);
Ok(fingerprint)
}
fn fingerprint_token(fingerprint: &[u8; MODEL_FINGERPRINT_BYTES]) -> String {
let mut token = String::with_capacity(MODEL_FINGERPRINT_BYTES * 2);
for byte in fingerprint {
write!(&mut token, "{byte:02x}").expect("writing to a String cannot fail");
}
token
}
struct ChineseClip {
model: ChineseClipModel,
tokenizer: Tokenizer,
device: Device,
}
impl ChineseClip {
fn load(model_path: &Path, vocabulary_path: &Path) -> Result<Self> {
let device = Device::Cpu;
// SAFETY: the mapped model is immutable for the process lifetime. The
// documentation requires stopping Pixhelf before replacing it.
let variables =
unsafe { VarBuilder::from_mmaped_safetensors(&[model_path], DType::F32, &device)? };
let config = ChineseClipConfig::clip_vit_base_patch16();
let model = ChineseClipModel::new(variables, &config)?;
let tokenizer = load_tokenizer(vocabulary_path)?;
Ok(Self {
model,
tokenizer,
device,
})
}
}
impl TextImageModel for ChineseClip {
fn embed_thumbnail(&mut self, thumbnail: &Path) -> Result<TextSearchEmbedding> {
let input = preprocess_thumbnail(thumbnail)?;
let values = self
.model
.get_image_features(&input)?
.i(0)?
.to_vec1::<f32>()?;
TextSearchEmbedding::quantize(&values)
}
fn embed_text(&mut self, text: &str) -> Result<TextSearchEmbedding> {
let encoding = self
.tokenizer
.encode(text, true)
.map_err(|error| anyhow!("cannot tokenize text query: {error}"))?;
if encoding.is_empty() {
bail!("text query did not produce any tokens");
}
let input_ids = Tensor::new(encoding.get_ids(), &self.device)?.unsqueeze(0)?;
let token_type_ids = Tensor::new(encoding.get_type_ids(), &self.device)?.unsqueeze(0)?;
let attention_mask =
Tensor::new(encoding.get_attention_mask(), &self.device)?.unsqueeze(0)?;
let values = self
.model
.get_text_features(&input_ids, Some(&token_type_ids), Some(&attention_mask))?
.i(0)?
.to_vec1::<f32>()?;
TextSearchEmbedding::quantize(&values)
}
}
fn load_tokenizer(vocabulary_path: &Path) -> Result<Tokenizer> {
let vocabulary = vocabulary_path.to_str().with_context(|| {
format!(
"vocabulary path is not UTF-8: {}",
vocabulary_path.display()
)
})?;
let model = WordPiece::from_file(vocabulary)
.unk_token("[UNK]".to_owned())
.build()
.map_err(|error| anyhow!("cannot load Chinese-CLIP vocabulary: {error}"))?;
let mut tokenizer = Tokenizer::new(model);
let cls_id = tokenizer
.token_to_id("[CLS]")
.context("Chinese-CLIP vocabulary is missing [CLS]")?;
let sep_id = tokenizer
.token_to_id("[SEP]")
.context("Chinese-CLIP vocabulary is missing [SEP]")?;
tokenizer.with_normalizer(Some(BertNormalizer::default()));
tokenizer.with_pre_tokenizer(Some(BertPreTokenizer));
tokenizer.with_post_processor(Some(BertProcessing::new(
("[SEP]".to_owned(), sep_id),
("[CLS]".to_owned(), cls_id),
)));
tokenizer
.with_truncation(Some(TruncationParams {
max_length: TEXT_CONTEXT_LENGTH,
..Default::default()
}))
.map_err(|error| anyhow!("cannot configure Chinese-CLIP tokenizer: {error}"))?;
Ok(tokenizer)
}
fn preprocess_thumbnail(path: &Path) -> Result<Tensor> {
let image = ImageReader::open(path)
.with_context(|| format!("cannot open thumbnail {}", path.display()))?
.with_guessed_format()?
.decode()
.with_context(|| format!("cannot decode thumbnail {}", path.display()))?;
preprocess_image(&image)
}
fn preprocess_image(image: &DynamicImage) -> Result<Tensor> {
let (width, height) = image.dimensions();
if width == 0 || height == 0 {
bail!("cannot embed an empty image");
}
// The official ChineseCLIPFeatureExtractor resizes directly to 224x224
// and does not centre-crop this checkpoint.
let pixels = image
.resize_exact(MODEL_EDGE, MODEL_EDGE, FilterType::CatmullRom)
.to_rgb8();
const MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73];
const STANDARD_DEVIATION: [f32; 3] = [0.268_629_54, 0.261_302_6, 0.275_777_1];
let mut values = Vec::with_capacity((MODEL_EDGE * MODEL_EDGE * 3) as usize);
for channel in 0..3 {
values.extend(pixels.pixels().map(|pixel| {
(f32::from(pixel.0[channel]) / 255.0 - MEAN[channel]) / STANDARD_DEVIATION[channel]
}));
}
Ok(Tensor::from_vec(
values,
(1, 3, MODEL_EDGE as usize, MODEL_EDGE as usize),
&Device::Cpu,
)?)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicUsize;
struct ColourAndWordModel {
image_calls: Arc<AtomicUsize>,
text_calls: Arc<AtomicUsize>,
}
impl TextImageModel for ColourAndWordModel {
fn embed_thumbnail(&mut self, thumbnail: &Path) -> Result<TextSearchEmbedding> {
self.image_calls.fetch_add(1, Ordering::Relaxed);
let pixel = image::open(thumbnail)?.to_rgb8().get_pixel(0, 0).0;
let axis = if pixel[0] >= pixel[2] { 0 } else { 1 };
Ok(TextSearchEmbedding::for_test([(axis, 100)]))
}
fn embed_text(&mut self, text: &str) -> Result<TextSearchEmbedding> {
self.text_calls.fetch_add(1, Ordering::Relaxed);
let axis = if text.contains('红') || text.contains("red") {
0
} else {
1
};
Ok(TextSearchEmbedding::for_test([(axis, 100)]))
}
}
#[test]
fn quantised_cross_modal_cosine_is_stable() {
let query = TextSearchEmbedding::for_test([(0, 100), (1, 50)]);
let same = TextSearchEmbedding::for_test([(0, 100), (1, 50)]);
let other = TextSearchEmbedding::for_test([(2, 100)]);
assert!((text_image_similarity(&query, &same) - 1.0).abs() < 1e-5);
assert!(text_image_similarity(&query, &other).abs() < 1e-5);
}
#[test]
fn sidecar_round_trips_and_rejects_other_models() {
let temp = tempfile::tempdir().unwrap();
let path = embedding_path(&temp.path().join("image.webp"));
let embedding = TextSearchEmbedding::for_test([(0, 100), (42, -70)]);
let fingerprint = [7_u8; MODEL_FINGERPRINT_BYTES];
write_embedding(&path, &fingerprint, &embedding).unwrap();
let loaded = read_embedding(&path, &fingerprint).unwrap();
assert!((text_image_similarity(&embedding, &loaded) - 1.0).abs() < 1e-5);
assert!(read_embedding(&path, &[8; MODEL_FINGERPRINT_BYTES]).is_err());
}
#[tokio::test]
async fn worker_persists_images_and_caches_text_queries() {
let temp = tempfile::tempdir().unwrap();
let thumbnail = temp.path().join("image.png");
image::RgbImage::from_pixel(8, 8, image::Rgb([240, 20, 10]))
.save(&thumbnail)
.unwrap();
let image_calls = Arc::new(AtomicUsize::new(0));
let text_calls = Arc::new(AtomicUsize::new(0));
let index = TextSearchIndex::with_test_model(
Box::new(ColourAndWordModel {
image_calls: Arc::clone(&image_calls),
text_calls: Arc::clone(&text_calls),
}),
[9; MODEL_FINGERPRINT_BYTES],
);
index.reconcile(&HashSet::from(["image"]));
index.enqueue("image", &thumbnail, false);
index.start_worker();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if index.status().background_complete {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
})
.await
.expect("text-search worker timed out");
let image = index.embedding("image").unwrap();
let first = index.embed_query("红色图片").await.unwrap();
let repeated = index.embed_query("红色图片").await.unwrap();
assert!(text_image_similarity(&first, &image) > 0.999);
assert!(text_image_similarity(&repeated, &image) > 0.999);
assert_eq!(image_calls.load(Ordering::Relaxed), 1);
assert_eq!(text_calls.load(Ordering::Relaxed), 1);
assert!(embedding_path(&thumbnail).is_file());
}
#[test]
fn tokenizer_uses_chinese_bert_special_tokens() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("vocab.txt");
let mut vocabulary = vec!["[PAD]".to_owned()];
vocabulary.extend((0..99).map(|index| format!("[unused{index}]")));
vocabulary.extend(
["[UNK]", "[CLS]", "[SEP]", "[MASK]", "海", "边", "日", "落"].map(str::to_owned),
);
fs::write(&path, format!("{}\n", vocabulary.join("\n"))).unwrap();
let tokenizer = load_tokenizer(&path).unwrap();
let encoding = tokenizer.encode("海边日落", true).unwrap();
assert_eq!(encoding.get_ids().first(), Some(&101));
assert_eq!(encoding.get_ids().last(), Some(&102));
assert!(encoding.len() >= 6);
}
#[test]
fn preprocessing_matches_official_chinese_clip_shape() {
let image = DynamicImage::ImageRgb8(image::RgbImage::from_fn(400, 200, |x, y| {
image::Rgb([(x % 255) as u8, (y % 255) as u8, 128])
}));
let tensor = preprocess_image(&image).unwrap();
assert_eq!(tensor.dims(), &[1, 3, 224, 224]);
assert!(tensor.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0].is_finite());
}
#[test]
fn official_model_runs_when_supplied_by_the_test_environment() {
let (Some(model_path), Some(vocabulary_path)) = (
std::env::var_os("PIXHELF_CHINESE_CLIP_MODEL"),
std::env::var_os("PIXHELF_CHINESE_CLIP_VOCAB"),
) else {
return;
};
let temp = tempfile::tempdir().unwrap();
let red = temp.path().join("red.png");
let blue = temp.path().join("blue.png");
image::RgbImage::from_pixel(224, 224, image::Rgb([235, 30, 25]))
.save(&red)
.unwrap();
image::RgbImage::from_pixel(224, 224, image::Rgb([25, 40, 235]))
.save(&blue)
.unwrap();
let mut model =
ChineseClip::load(Path::new(&model_path), Path::new(&vocabulary_path)).unwrap();
let query = model.embed_text("一张红色的图片").unwrap();
let red = model.embed_thumbnail(&red).unwrap();
let blue = model.embed_thumbnail(&blue).unwrap();
let red_score = text_image_similarity(&query, &red);
let blue_score = text_image_similarity(&query, &blue);
eprintln!("official Chinese-CLIP scores: red={red_score:.4}, blue={blue_score:.4}");
assert!(red_score > blue_score);
}
}
src/thumbs.rs
+249 -2
View File
@@ -22,7 +22,21 @@ use tokio::sync::Notify;
use tracing::{info, warn};
use walkdir::WalkDir;
use crate::gallery::ImageRecord;
use crate::{
gallery::ImageRecord,
semantic::{
SemanticEmbedding, SemanticIndex, SemanticStatus, is_semantic_sidecar,
remove_embedding_for_thumbnail,
},
similarity::{
cache_signature_for_thumbnail, has_signature_sidecar, is_signature_sidecar,
remove_signature_for_thumbnail, signature_for_thumbnail,
},
text_search::{
TextSearchEmbedding, TextSearchIndex, TextSearchStatus, is_text_search_sidecar,
remove_embedding_for_thumbnail as remove_text_search_embedding,
},
};
const THUMBNAIL_EDGE: u32 = 720;
const WEBP_QUALITY: f32 = 82.0;
@@ -37,6 +51,9 @@ pub struct ThumbnailManager {
entries: RwLock<HashMap<String, Arc<ThumbEntry>>>,
queue: JobQueue,
initial_ready: AtomicBool,
similarity_warmup_running: AtomicBool,
semantic: Option<Arc<SemanticIndex>>,
text_search: Option<Arc<TextSearchIndex>>,
}
#[derive(Debug)]
@@ -109,10 +126,21 @@ pub struct ThumbnailStatus {
pub failed: usize,
pub initial_batch_ready: bool,
pub background_complete: bool,
pub semantic: SemanticStatus,
pub text_search: TextSearchStatus,
}
impl ThumbnailManager {
#[cfg(test)]
pub fn new(cache_dir: PathBuf) -> Result<Arc<Self>> {
Self::new_with_indexes(cache_dir, None, None)
}
pub fn new_with_indexes(
cache_dir: PathBuf,
semantic: Option<Arc<SemanticIndex>>,
text_search: Option<Arc<TextSearchIndex>>,
) -> Result<Arc<Self>> {
let cache_root = cache_dir.join(CACHE_VERSION);
fs::create_dir_all(&cache_root)
.with_context(|| format!("cannot create thumbnail cache: {}", cache_root.display()))?;
@@ -132,6 +160,9 @@ impl ThumbnailManager {
entries: RwLock::new(HashMap::new()),
queue: JobQueue::default(),
initial_ready: AtomicBool::new(false),
similarity_warmup_running: AtomicBool::new(false),
semantic,
text_search,
}))
}
@@ -179,6 +210,31 @@ impl ThumbnailManager {
for id in missing {
self.queue.push_background(id);
}
if self.semantic.is_some() || self.text_search.is_some() {
if let Some(semantic) = &self.semantic {
semantic.reconcile(&current_ids);
}
if let Some(text_search) = &self.text_search {
text_search.reconcile(&current_ids);
}
let ready = {
let entries = read_lock(&self.entries);
entries
.values()
.filter(|entry| matches!(*mutex_lock(&entry.state), ThumbState::Ready))
.filter(|entry| valid_cache_file(&entry.cache_path))
.map(|entry| (entry.record.id.clone(), entry.cache_path.clone()))
.collect::<Vec<_>>()
};
for (id, thumbnail) in ready {
if let Some(semantic) = &self.semantic {
semantic.enqueue(&id, &thumbnail, false);
}
if let Some(text_search) = &self.text_search {
text_search.enqueue(&id, &thumbnail, false);
}
}
}
}
pub fn start_workers(self: &Arc<Self>, count: usize) {
@@ -190,6 +246,53 @@ impl ThumbnailManager {
}
}
/// Build descriptors for thumbnails created by an older Pixhelf release.
/// New thumbnails write their descriptor during generation, so this is a
/// one-time, low-priority migration after an upgrade.
pub fn start_similarity_warmup(self: &Arc<Self>) {
if self.similarity_warmup_running.swap(true, Ordering::AcqRel) {
return;
}
let manager = Arc::clone(self);
tokio::spawn(async move {
let pending = {
let entries = read_lock(&manager.entries);
entries
.values()
.filter(|entry| matches!(*mutex_lock(&entry.state), ThumbState::Ready))
.filter(|entry| valid_cache_file(&entry.cache_path))
.filter(|entry| !has_signature_sidecar(&entry.cache_path))
.map(|entry| (Arc::clone(&entry.record), entry.cache_path.clone()))
.collect::<Vec<_>>()
};
let total = pending.len();
let result = tokio::task::spawn_blocking(move || {
pending
.into_iter()
.filter(|(record, thumbnail)| {
if let Err(error) = signature_for_thumbnail(record, thumbnail) {
warn!(image = %record.relative_path, %error, "cannot warm image similarity descriptor");
true
} else {
false
}
})
.count()
})
.await;
match result {
Ok(failures) if total > 0 => {
info!(total, failures, "image similarity descriptors are ready");
}
Ok(_) => {}
Err(error) => warn!(%error, "image similarity descriptor warmup stopped"),
}
manager
.similarity_warmup_running
.store(false, Ordering::Release);
});
}
pub async fn prepare_initial(&self, ids: &[String]) {
for id in ids {
self.promote(id);
@@ -240,9 +343,62 @@ impl ThumbnailManager {
failed,
initial_batch_ready: self.initial_ready.load(Ordering::Acquire),
background_complete: ready + failed == total && processing == 0,
semantic: self
.semantic
.as_ref()
.map_or_else(SemanticStatus::disabled, |semantic| semantic.status()),
text_search: self
.text_search
.as_ref()
.map_or_else(TextSearchStatus::disabled, |index| index.status()),
}
}
pub async fn ensure_semantic_embedding(
&self,
id: &str,
thumbnail: &Path,
) -> Option<Arc<SemanticEmbedding>> {
let semantic = self.semantic.as_ref()?;
semantic.ensure_embedding(id, thumbnail).await
}
pub fn semantic_embedding(&self, id: &str) -> Option<Arc<SemanticEmbedding>> {
self.semantic.as_ref()?.embedding(id)
}
pub fn semantic_cache_token(&self) -> String {
self.semantic
.as_ref()
.map_or_else(|| "off".to_owned(), |semantic| semantic.cache_token())
}
pub async fn text_query_embedding(
&self,
query: &str,
) -> Result<Option<Arc<TextSearchEmbedding>>> {
let Some(index) = &self.text_search else {
return Ok(None);
};
index.embed_query(query).await.map(Some)
}
pub fn text_search_embedding(&self, id: &str) -> Option<Arc<TextSearchEmbedding>> {
self.text_search.as_ref()?.embedding(id)
}
pub fn text_search_cache_token(&self) -> String {
self.text_search
.as_ref()
.map_or_else(|| "off".to_owned(), |index| index.cache_token())
}
pub fn ready_path(&self, id: &str) -> Option<PathBuf> {
let entry = read_lock(&self.entries).get(id).cloned()?;
let ready = matches!(*mutex_lock(&entry.state), ThumbState::Ready);
(ready && valid_cache_file(&entry.cache_path)).then(|| entry.cache_path.clone())
}
pub async fn cleanup_stale(&self) {
let valid_ids: HashSet<String> = read_lock(&self.entries).keys().cloned().collect();
let root = self.cache_root.clone();
@@ -354,6 +510,12 @@ impl ThumbnailManager {
}
};
if ready {
if let Some(semantic) = &self.semantic {
semantic.enqueue(&id, &entry.cache_path, false);
}
if let Some(text_search) = &self.text_search {
text_search.enqueue(&id, &entry.cache_path, false);
}
entry.notify.notify_waiters();
} else if removed {
let entries = read_lock(&self.entries);
@@ -363,6 +525,21 @@ impl ThumbnailManager {
{
warn!(path = %entry.cache_path.display(), %error, "cannot remove stale thumbnail");
}
if !entries.contains_key(&id)
&& let Err(error) = remove_signature_for_thumbnail(&entry.cache_path)
{
warn!(path = %entry.cache_path.display(), %error, "cannot remove stale similarity descriptor");
}
if !entries.contains_key(&id)
&& let Err(error) = remove_embedding_for_thumbnail(&entry.cache_path)
{
warn!(path = %entry.cache_path.display(), %error, "cannot remove stale semantic embedding");
}
if !entries.contains_key(&id)
&& let Err(error) = remove_text_search_embedding(&entry.cache_path)
{
warn!(path = %entry.cache_path.display(), %error, "cannot remove stale text-search embedding");
}
}
}
Err(error) => {
@@ -516,12 +693,24 @@ fn generate_thumbnail(record: &ImageRecord, output: &Path) -> Result<()> {
};
write_webp(&image, output, WEBP_QUALITY)
.with_context(|| format!("cannot store thumbnail: {}", output.display()))?;
if let Err(error) = cache_signature_for_thumbnail(record, &image, output) {
warn!(path = %output.display(), %error, "cannot cache image similarity descriptor");
}
if let Err(error) = record.ensure_source_is_current() {
if let Err(remove_error) = fs::remove_file(output)
&& remove_error.kind() != std::io::ErrorKind::NotFound
{
warn!(path = %output.display(), %remove_error, "cannot remove stale thumbnail");
}
if let Err(remove_error) = remove_signature_for_thumbnail(output) {
warn!(path = %output.display(), %remove_error, "cannot remove stale similarity descriptor");
}
if let Err(remove_error) = remove_embedding_for_thumbnail(output) {
warn!(path = %output.display(), %remove_error, "cannot remove stale semantic embedding");
}
if let Err(remove_error) = remove_text_search_embedding(output) {
warn!(path = %output.display(), %remove_error, "cannot remove stale text-search embedding");
}
return Err(error);
}
Ok(())
@@ -573,7 +762,11 @@ fn write_webp(image: &DynamicImage, output: &Path, quality: f32) -> Result<()> {
fn cleanup_cache(root: &Path, valid_ids: &HashSet<String>) {
visit_cache_files(root, |path| {
if path.extension().and_then(|ext| ext.to_str()) != Some("webp") {
if path.extension().and_then(|ext| ext.to_str()) != Some("webp")
&& !is_signature_sidecar(path)
&& !is_semantic_sidecar(path)
&& !is_text_search_sidecar(path)
{
return;
}
let stem = path
@@ -627,6 +820,7 @@ mod tests {
generate_thumbnail(&index.images[0], &output).unwrap();
let generated = image::open(output).unwrap();
assert_eq!(generated.dimensions(), (720, 360));
assert!(temp.path().join("thumbnail.sim2").is_file());
}
#[test]
@@ -648,19 +842,72 @@ mod tests {
let root = temp.path();
let id = "abcdef";
let thumbnail = shard_path(root, id);
let signature = thumbnail.with_extension("sim2");
let semantic = thumbnail.with_extension("dino2");
let text_search = thumbnail.with_extension("cnclip");
let old_compact = thumbnail.with_file_name(format!("{id}-360.webp"));
let stale = shard_path(root, "stale");
let stale_signature = stale.with_extension("sim2");
let stale_semantic = stale.with_extension("dino2");
let stale_text_search = stale.with_extension("cnclip");
fs::create_dir_all(thumbnail.parent().unwrap()).unwrap();
fs::create_dir_all(stale.parent().unwrap()).unwrap();
fs::write(&thumbnail, [1; 32]).unwrap();
fs::write(&signature, [1; 32]).unwrap();
fs::write(&semantic, [1; 32]).unwrap();
fs::write(&text_search, [1; 32]).unwrap();
fs::write(&old_compact, [1; 32]).unwrap();
fs::write(&stale, [1; 32]).unwrap();
fs::write(&stale_signature, [1; 32]).unwrap();
fs::write(&stale_semantic, [1; 32]).unwrap();
fs::write(&stale_text_search, [1; 32]).unwrap();
cleanup_cache(root, &HashSet::from([id.to_owned()]));
assert!(thumbnail.exists());
assert!(signature.exists());
assert!(semantic.exists());
assert!(text_search.exists());
assert!(!old_compact.exists());
assert!(!stale.exists());
assert!(!stale_signature.exists());
assert!(!stale_semantic.exists());
assert!(!stale_text_search.exists());
}
#[tokio::test]
async fn warmup_migrates_cached_thumbnails_without_descriptors() {
let temp = tempfile::tempdir().unwrap();
let gallery = temp.path().join("gallery");
fs::create_dir(&gallery).unwrap();
let source = gallery.join("warmup.png");
image::RgbImage::from_fn(48, 32, |x, y| {
image::Rgb([(x * 4) as u8, (y * 6) as u8, ((x + y) * 3) as u8])
})
.save(&source)
.unwrap();
let index = scan_gallery(&gallery, None).unwrap();
let manager = ThumbnailManager::new(temp.path().join("cache")).unwrap();
let cached = manager.cache_path(&index.images[0].id);
fs::create_dir_all(cached.parent().unwrap()).unwrap();
write_webp(&image::open(&source).unwrap(), &cached, WEBP_QUALITY).unwrap();
assert!(valid_cache_file(&cached));
assert!(!has_signature_sidecar(&cached));
manager.reconcile(&index.images);
manager.start_similarity_warmup();
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if has_signature_sidecar(&cached)
&& !manager.similarity_warmup_running.load(Ordering::Acquire)
{
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("similarity descriptor warmup timed out");
}
#[test]
src/web.rs
+878 -6
View File
@@ -1,4 +1,8 @@
use std::{cmp::Reverse, sync::Arc};
use std::{
cmp::Reverse,
collections::HashMap,
sync::{Arc, Mutex, OnceLock},
};
use axum::{
Json, Router,
@@ -9,7 +13,7 @@ use axum::{
routing::get,
};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tokio::sync::{RwLock, Semaphore};
use tower::ServiceExt;
use tower_http::{
compression::{CompressionLayer, CompressionLevel},
@@ -21,6 +25,12 @@ use tracing::error;
use crate::{
gallery::{Album, GalleryIndex, ImageRecord},
semantic::{SemanticEmbedding, semantic_similarity},
similarity::{
DiversityFingerprint, MIN_SIMILARITY_SCORE, redundancy_score, signature_for_thumbnail,
similarity_score,
},
text_search::{TextSearchEmbedding, text_image_similarity},
thumbs::{ThumbnailError, ThumbnailManager, ThumbnailStatus},
};
@@ -39,8 +49,86 @@ const APP_CSS: &[u8] = include_bytes!(concat!(
const ASSET_VERSION: &str = env!("PIXHELF_ASSET_VERSION");
const DEFAULT_PAGE_SIZE: usize = 60;
const MAX_PAGE_SIZE: usize = 200;
const DIVERSIFIED_SIMILARITY_RESULTS: usize = 120;
const REDUNDANCY_PENALTY_START: f32 = 0.86;
const MAX_REDUNDANCY_PENALTY: f32 = 0.08;
const SEMANTIC_WEIGHT: f32 = 0.70;
// DINOv2 CLS vectors are zero-centred: unrelated images cluster around zero,
// so raw cosine values are not probabilities. Preserve the raw gate for
// precision, then calibrate the useful -0.10..0.20 range before blending it
// with the model-free visual score.
const MIN_SEMANTIC_COSINE: f32 = 0.085;
const SEMANTIC_CALIBRATION_OFFSET: f32 = 0.10;
const SEMANTIC_CALIBRATION_RANGE: f32 = 0.30;
const MIN_HYBRID_SCORE: f32 = 0.48;
const SIMILARITY_RANKING_CACHE_LIMIT: usize = 8;
const TEXT_SEARCH_RANKING_CACHE_LIMIT: usize = 16;
const TEXT_SEARCH_RESULT_LIMIT: usize = 600;
const FILENAME_MATCH_BOOST: f32 = 3.0;
const MAX_QUERY_VALUE_BYTES: usize = 4096;
static SIMILARITY_LIMIT: OnceLock<Arc<Semaphore>> = OnceLock::new();
type SimilarityRanking = Arc<Vec<Arc<ImageRecord>>>;
type SimilarityRankingCache = Mutex<HashMap<String, SimilarityRanking>>;
static SIMILARITY_RANKING_CACHE: OnceLock<SimilarityRankingCache> = OnceLock::new();
type TextSearchRanking = Arc<Vec<Arc<ImageRecord>>>;
type TextSearchRankingCache = Mutex<HashMap<String, TextSearchRanking>>;
static TEXT_SEARCH_RANKING_CACHE: OnceLock<TextSearchRankingCache> = OnceLock::new();
fn similarity_limit() -> Arc<Semaphore> {
Arc::clone(SIMILARITY_LIMIT.get_or_init(|| Arc::new(Semaphore::new(1))))
}
fn similarity_ranking_cache() -> &'static SimilarityRankingCache {
SIMILARITY_RANKING_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn cached_similarity_ranking(key: &str) -> Option<SimilarityRanking> {
similarity_ranking_cache()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.get(key)
.cloned()
}
fn cache_similarity_ranking(key: String, ranking: SimilarityRanking) {
let mut cache = similarity_ranking_cache()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !cache.contains_key(&key)
&& cache.len() >= SIMILARITY_RANKING_CACHE_LIMIT
&& let Some(evicted) = cache.keys().next().cloned()
{
cache.remove(&evicted);
}
cache.insert(key, ranking);
}
fn text_search_ranking_cache() -> &'static TextSearchRankingCache {
TEXT_SEARCH_RANKING_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn cached_text_search_ranking(key: &str) -> Option<TextSearchRanking> {
text_search_ranking_cache()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.get(key)
.cloned()
}
fn cache_text_search_ranking(key: String, ranking: TextSearchRanking) {
let mut cache = text_search_ranking_cache()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !cache.contains_key(&key)
&& cache.len() >= TEXT_SEARCH_RANKING_CACHE_LIMIT
&& let Some(evicted) = cache.keys().next().cloned()
{
cache.remove(&evicted);
}
cache.insert(key, ranking);
}
#[derive(Clone)]
pub struct AppState {
pub index: Arc<RwLock<Arc<GalleryIndex>>>,
@@ -59,6 +147,30 @@ struct ImagesQuery {
limit: Option<usize>,
}
#[derive(Debug, Default, Deserialize)]
struct SimilarImagesQuery {
offset: Option<usize>,
limit: Option<usize>,
}
#[derive(Debug, Clone, Copy)]
struct SimilarImagesRequest {
offset: usize,
limit: usize,
}
impl SimilarImagesQuery {
fn normalize(self) -> SimilarImagesRequest {
SimilarImagesRequest {
offset: self.offset.unwrap_or(0),
limit: self
.limit
.unwrap_or(DEFAULT_PAGE_SIZE)
.clamp(1, MAX_PAGE_SIZE),
}
}
}
#[derive(Debug)]
struct ImageRequest {
album: Option<String>,
@@ -172,6 +284,7 @@ pub fn router(state: AppState) -> Router {
.route("/api/health", get(health))
.route("/api/gallery", get(gallery_summary))
.route("/api/images", get(images))
.route("/api/images/{id}/similar", get(similar_images))
.route("/api/status", get(thumbnail_status))
.route("/api/images/{id}/thumbnail", get(thumbnail))
.route("/api/images/{id}/original", get(original))
@@ -230,12 +343,86 @@ async fn images(
Ok(query) => query,
Err(status) => return status.into_response(),
};
let index = state.index.read().await;
let etag = images_etag(&index.revision, &query);
let index = {
let index = state.index.read().await;
Arc::clone(&index)
};
let status = state.thumbnails.status();
let text_search_active = query.search.is_some()
&& status.background_complete
&& status.text_search.enabled
&& status.text_search.background_complete
&& status.text_search.ready > 0;
let text_search_token = text_search_active.then(|| state.thumbnails.text_search_cache_token());
let etag = images_etag(&index.revision, &query, text_search_token.as_deref());
if is_not_modified(&headers, &etag) {
return not_modified(&etag, "private, no-cache");
}
if text_search_active {
let search = query.search.as_deref().expect("search was checked above");
let ranking_key = text_search_ranking_key(
&index.revision,
query.album.as_deref(),
search,
text_search_token.as_deref().unwrap_or("off"),
);
if let Some(ranking) = cached_text_search_ranking(&ranking_key) {
return with_cache_headers(
Json(text_search_image_page(&ranking, &query)).into_response(),
&etag,
"private, no-cache",
);
}
match state.thumbnails.text_query_embedding(search).await {
Ok(Some(query_embedding)) => {
let candidates = index
.images
.iter()
.filter(|record| {
query
.album
.as_deref()
.is_none_or(|album| record.album == album)
})
.filter_map(|record| {
let embedding = state.thumbnails.text_search_embedding(&record.id);
let filename_match = record.search_key.contains(search);
(embedding.is_some() || filename_match)
.then(|| (Arc::clone(record), embedding, filename_match))
})
.collect::<Vec<_>>();
let ranking = match tokio::task::spawn_blocking(move || {
rank_text_search(query_embedding, candidates)
})
.await
{
Ok(ranking) => Arc::new(ranking),
Err(error) => {
error!(%error, "natural-language image ranking task stopped");
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
};
cache_text_search_ranking(ranking_key, Arc::clone(&ranking));
return with_cache_headers(
Json(text_search_image_page(&ranking, &query)).into_response(),
&etag,
"private, no-cache",
);
}
Ok(None) => {}
Err(error) => {
error!(query = %search, %error, "cannot encode natural-language image query; using filename fallback");
let mut response = Json(image_page(&index, &query)).into_response();
response
.headers_mut()
.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
return response;
}
}
}
with_cache_headers(
Json(image_page(&index, &query)).into_response(),
&etag,
@@ -243,6 +430,315 @@ async fn images(
)
}
async fn similar_images(
State(state): State<AppState>,
Path(id): Path<String>,
Query(query): Query<SimilarImagesQuery>,
headers: HeaderMap,
) -> Response {
let query = query.normalize();
let index = {
let index = state.index.read().await;
Arc::clone(&index)
};
let Some(source) = index.image(&id).cloned() else {
return StatusCode::NOT_FOUND.into_response();
};
let source_thumbnail = match state.thumbnails.ensure_ready(&id).await {
Ok(path) => path,
Err(ThumbnailError::NotFound) => return StatusCode::NOT_FOUND.into_response(),
Err(error) => {
error!(%id, %error, "cannot prepare source thumbnail for similarity search");
return (StatusCode::INTERNAL_SERVER_ERROR, "无法读取查询图片").into_response();
}
};
let source_semantic = state
.thumbnails
.ensure_semantic_embedding(&id, &source_thumbnail)
.await;
let initial_status = state.thumbnails.status();
let initial_semantic_token = state.thumbnails.semantic_cache_token();
let initial_ranking_key = similar_ranking_key(
&index.revision,
&id,
initial_status.ready,
initial_status.failed,
&initial_semantic_token,
);
let initial_etag = similar_images_etag(
&index.revision,
&id,
initial_status.ready,
initial_status.failed,
&initial_semantic_token,
query,
);
if is_not_modified(&headers, &initial_etag) {
return not_modified(&initial_etag, "private, no-cache");
}
if let Some(ranking) = cached_similarity_ranking(&initial_ranking_key) {
return with_cache_headers(
Json(similar_image_page(&ranking, query)).into_response(),
&initial_etag,
"private, no-cache",
);
}
let permit = match similarity_limit().acquire_owned().await {
Ok(permit) => permit,
Err(error) => {
error!(%id, %error, "similarity search queue stopped");
return StatusCode::SERVICE_UNAVAILABLE.into_response();
}
};
// Thumbnail generation can finish while this request waits for the single
// ranking worker. Refresh the key and conditional response before doing
// any expensive descriptor work.
let thumbnail_status = state.thumbnails.status();
let semantic_token = state.thumbnails.semantic_cache_token();
let ranking_key = similar_ranking_key(
&index.revision,
&id,
thumbnail_status.ready,
thumbnail_status.failed,
&semantic_token,
);
let etag = similar_images_etag(
&index.revision,
&id,
thumbnail_status.ready,
thumbnail_status.failed,
&semantic_token,
query,
);
if is_not_modified(&headers, &etag) {
return not_modified(&etag, "private, no-cache");
}
if let Some(ranking) = cached_similarity_ranking(&ranking_key) {
return with_cache_headers(
Json(similar_image_page(&ranking, query)).into_response(),
&etag,
"private, no-cache",
);
}
// Keep results stable while the one-time semantic index is still being
// built. The frontend refreshes an active recommendation when indexing
// completes; until then every candidate uses the dependable local visual
// descriptor instead of mixing indexed and not-yet-indexed candidates.
let semantic_active = thumbnail_status.semantic.enabled
&& thumbnail_status.semantic.background_complete
&& source_semantic.is_some();
let source_semantic = semantic_active.then_some(source_semantic).flatten();
let candidates = index
.images
.iter()
.filter(|candidate| candidate.id != source.id)
.filter_map(|candidate| {
state.thumbnails.ready_path(&candidate.id).map(|path| {
let semantic = semantic_active
.then(|| state.thumbnails.semantic_embedding(&candidate.id))
.flatten();
(Arc::clone(candidate), path, semantic)
})
})
.collect::<Vec<_>>();
let ranking = match tokio::task::spawn_blocking(move || {
let _permit = permit;
rank_similar_images(source, source_thumbnail, source_semantic, candidates)
})
.await
{
Ok(Ok(ranking)) => Arc::new(ranking),
Ok(Err(error)) => {
error!(%id, %error, "cannot calculate similar images");
return (StatusCode::INTERNAL_SERVER_ERROR, "无法计算相似图片").into_response();
}
Err(error) => {
error!(%id, %error, "similar image task stopped");
return (StatusCode::INTERNAL_SERVER_ERROR, "无法计算相似图片").into_response();
}
};
cache_similarity_ranking(ranking_key, Arc::clone(&ranking));
let page = similar_image_page(&ranking, query);
with_cache_headers(Json(page).into_response(), &etag, "private, no-cache")
}
struct RankedSimilarity {
record: Arc<ImageRecord>,
score: f32,
fingerprint: DiversityFingerprint,
}
#[derive(Clone, Copy)]
struct SimilarityScore {
relevance: f32,
included: bool,
}
fn rank_similar_images(
source: Arc<ImageRecord>,
source_thumbnail: std::path::PathBuf,
source_semantic: Option<Arc<SemanticEmbedding>>,
candidates: Vec<(
Arc<ImageRecord>,
std::path::PathBuf,
Option<Arc<SemanticEmbedding>>,
)>,
) -> anyhow::Result<Vec<Arc<ImageRecord>>> {
let source_signature = signature_for_thumbnail(source.as_ref(), &source_thumbnail)?;
let worker_count = std::thread::available_parallelism()
.map(|parallelism| parallelism.get().clamp(2, 8))
.unwrap_or(2)
.min(candidates.len().max(1));
let chunk_size = candidates.len().div_ceil(worker_count).max(1);
let source_semantic = source_semantic.as_deref();
let mut ranked = std::thread::scope(|scope| -> anyhow::Result<Vec<_>> {
let workers = candidates
.chunks(chunk_size)
.map(|chunk| {
let source_id = source.id.as_str();
let source_signature = &source_signature;
scope.spawn(move || {
let mut matches = Vec::with_capacity(chunk.len());
for (candidate, thumbnail, candidate_semantic) in chunk {
if candidate.id == source_id {
continue;
}
// A malformed candidate should not make the whole gallery
// unusable. It is omitted until a rescan fixes or removes it.
let Ok(candidate_signature) =
signature_for_thumbnail(candidate.as_ref(), thumbnail)
else {
continue;
};
let visual_score = similarity_score(source_signature, &candidate_signature);
let semantic_score = source_semantic
.zip(candidate_semantic.as_deref())
.map(|(source, candidate)| semantic_similarity(source, candidate));
let score = hybrid_similarity_score(visual_score, semantic_score);
if !score.included {
continue;
}
matches.push(RankedSimilarity {
record: Arc::clone(candidate),
score: score.relevance,
fingerprint: candidate_signature.diversity_fingerprint(),
});
}
matches
})
})
.collect::<Vec<_>>();
let mut matches = Vec::with_capacity(candidates.len().saturating_sub(1));
for worker in workers {
matches.extend(
worker
.join()
.map_err(|_| anyhow::anyhow!("similarity worker stopped"))?,
);
}
Ok(matches)
})?;
ranked.sort_by(|left, right| {
right
.score
.total_cmp(&left.score)
.then_with(|| left.record.id.cmp(&right.record.id))
});
diversify_similarity_results(&mut ranked);
Ok(ranked
.into_iter()
.map(|candidate| candidate.record)
.collect())
}
fn hybrid_similarity_score(visual: f32, semantic: Option<f32>) -> SimilarityScore {
let Some(semantic) = semantic else {
return SimilarityScore {
relevance: visual,
included: visual >= MIN_SIMILARITY_SCORE,
};
};
let semantic_cosine = semantic.clamp(-1.0, 1.0);
let semantic_relevance = ((semantic_cosine + SEMANTIC_CALIBRATION_OFFSET)
/ SEMANTIC_CALIBRATION_RANGE)
.clamp(0.0, 1.0);
let relevance = semantic_relevance * SEMANTIC_WEIGHT + visual * (1.0 - SEMANTIC_WEIGHT);
SimilarityScore {
relevance,
included: visual >= MIN_SIMILARITY_SCORE
|| (semantic_cosine >= MIN_SEMANTIC_COSINE && relevance >= MIN_HYBRID_SCORE),
}
}
fn similar_image_page(ranked: &[Arc<ImageRecord>], query: SimilarImagesRequest) -> ImagesResponse {
let total = ranked.len();
let offset = query.offset.min(total);
let end = offset.saturating_add(query.limit).min(total);
let items = ranked[offset..end]
.iter()
.map(|record| ImageView::from(record.as_ref()))
.collect();
finish_page(items, total, offset, query.limit)
}
fn diversify_similarity_results(ranked: &mut Vec<RankedSimilarity>) {
let target = ranked.len().min(DIVERSIFIED_SIMILARITY_RESULTS);
if target < 2 {
return;
}
let mut remaining = std::mem::take(ranked)
.into_iter()
.map(|candidate| (candidate, 0.0_f32))
.collect::<Vec<_>>();
let mut selected = Vec::with_capacity(target);
while selected.len() < target {
let next = remaining
.iter()
.enumerate()
.max_by(
|(_, (left, left_redundancy)), (_, (right, right_redundancy))| {
diversity_rank(left.score, *left_redundancy)
.total_cmp(&diversity_rank(right.score, *right_redundancy))
.then_with(|| right.record.id.cmp(&left.record.id))
},
)
.map(|(index, _)| index)
.expect("similarity selection has remaining candidates");
let (candidate, _) = remaining.swap_remove(next);
let fingerprint = candidate.fingerprint;
selected.push(candidate);
for (remaining_candidate, maximum_redundancy) in &mut remaining {
*maximum_redundancy = maximum_redundancy.max(redundancy_score(
fingerprint,
remaining_candidate.fingerprint,
));
}
}
remaining.sort_by(|(left, _), (right, _)| {
right
.score
.total_cmp(&left.score)
.then_with(|| left.record.id.cmp(&right.record.id))
});
selected.extend(remaining.into_iter().map(|(candidate, _)| candidate));
*ranked = selected;
}
fn diversity_rank(relevance: f32, redundancy: f32) -> f32 {
let penalty = ((redundancy - REDUNDANCY_PENALTY_START) / (1.0 - REDUNDANCY_PENALTY_START))
.clamp(0.0, 1.0)
.powi(2)
* MAX_REDUNDANCY_PENALTY;
relevance - penalty
}
fn image_page(index: &GalleryIndex, query: &ImageRequest) -> ImagesResponse {
let album = query.album.as_deref();
let search = query.search.as_deref();
@@ -291,6 +787,54 @@ fn image_page(index: &GalleryIndex, query: &ImageRequest) -> ImagesResponse {
}
}
struct RankedTextSearch {
record: Arc<ImageRecord>,
score: f32,
}
fn rank_text_search(
query: Arc<TextSearchEmbedding>,
candidates: Vec<(Arc<ImageRecord>, Option<Arc<TextSearchEmbedding>>, bool)>,
) -> Vec<Arc<ImageRecord>> {
let mut ranked = candidates
.into_iter()
.map(|(record, embedding, filename_match)| {
let semantic = embedding
.as_deref()
.map_or(-1.0, |embedding| text_image_similarity(&query, embedding));
let score = semantic
+ if filename_match {
FILENAME_MATCH_BOOST
} else {
0.0
};
RankedTextSearch { record, score }
})
.collect::<Vec<_>>();
ranked.sort_by(|left, right| {
right
.score
.total_cmp(&left.score)
.then_with(|| left.record.id.cmp(&right.record.id))
});
ranked
.into_iter()
.take(TEXT_SEARCH_RESULT_LIMIT)
.map(|ranked| ranked.record)
.collect()
}
fn text_search_image_page(ranked: &[Arc<ImageRecord>], query: &ImageRequest) -> ImagesResponse {
let total = ranked.len();
let offset = query.offset.min(total);
let end = offset.saturating_add(query.limit).min(total);
let items = ranked[offset..end]
.iter()
.map(|record| ImageView::from(record.as_ref()))
.collect();
finish_page(items, total, offset, query.limit)
}
fn page_from_unfiltered_index(
index: &GalleryIndex,
reverse: bool,
@@ -475,21 +1019,71 @@ fn revision_etag(kind: &str, revision: &str) -> String {
format!("\"{kind}-{revision}\"")
}
fn images_etag(revision: &str, query: &ImageRequest) -> String {
fn images_etag(revision: &str, query: &ImageRequest, text_search_token: Option<&str>) -> String {
let mut hasher = blake3::Hasher::new();
hasher.update(b"pixhelf-images-query-v1\0");
hasher.update(b"pixhelf-images-query-v2\0");
hasher.update(query.album.as_deref().unwrap_or_default().as_bytes());
hasher.update(b"\0");
hasher.update(query.search.as_deref().unwrap_or_default().as_bytes());
hasher.update(b"\0");
hasher.update(&[query.sort as u8]);
hasher.update(query.seed.as_bytes());
hasher.update(b"\0");
hasher.update(text_search_token.unwrap_or("filename").as_bytes());
hasher.update(&query.offset.to_le_bytes());
hasher.update(&query.limit.to_le_bytes());
let query_hash = hasher.finalize().to_hex();
format!("\"images-{revision}-{}\"", &query_hash[..16])
}
fn text_search_ranking_key(
revision: &str,
album: Option<&str>,
search: &str,
text_search_token: &str,
) -> String {
let mut hasher = blake3::Hasher::new();
hasher.update(b"pixhelf-text-search-v1\0");
hasher.update(revision.as_bytes());
hasher.update(b"\0");
hasher.update(album.unwrap_or_default().as_bytes());
hasher.update(b"\0");
hasher.update(search.as_bytes());
hasher.update(b"\0");
hasher.update(text_search_token.as_bytes());
hasher.finalize().to_hex().to_string()
}
fn similar_images_etag(
revision: &str,
source_id: &str,
ready: usize,
failed: usize,
semantic_token: &str,
query: SimilarImagesRequest,
) -> String {
let mut hasher = blake3::Hasher::new();
hasher.update(b"pixhelf-similar-images-v3\0");
hasher.update(source_id.as_bytes());
hasher.update(&ready.to_le_bytes());
hasher.update(&failed.to_le_bytes());
hasher.update(semantic_token.as_bytes());
hasher.update(&query.offset.to_le_bytes());
hasher.update(&query.limit.to_le_bytes());
let query_hash = hasher.finalize().to_hex();
format!("\"similar-{revision}-{}\"", &query_hash[..16])
}
fn similar_ranking_key(
revision: &str,
source_id: &str,
ready: usize,
failed: usize,
semantic_token: &str,
) -> String {
format!("{revision}:{source_id}:{ready}:{failed}:{semantic_token}")
}
fn is_not_modified(headers: &HeaderMap, etag: &str) -> bool {
headers
.get(header::IF_NONE_MATCH)
@@ -631,6 +1225,24 @@ fn escape_script_json(json: String) -> String {
mod tests {
use super::*;
struct ColourTextSearchModel;
impl crate::text_search::TextImageModel for ColourTextSearchModel {
fn embed_thumbnail(
&mut self,
thumbnail: &std::path::Path,
) -> anyhow::Result<TextSearchEmbedding> {
let pixel = image::open(thumbnail)?.to_rgb8().get_pixel(0, 0).0;
let axis = if pixel[0] >= pixel[2] { 0 } else { 1 };
Ok(TextSearchEmbedding::for_test([(axis, 100)]))
}
fn embed_text(&mut self, text: &str) -> anyhow::Result<TextSearchEmbedding> {
let axis = if text.contains('红') { 0 } else { 1 };
Ok(TextSearchEmbedding::for_test([(axis, 100)]))
}
}
fn empty_router() -> Router {
let temp = tempfile::tempdir().unwrap();
let thumbnails = ThumbnailManager::new(temp.path().to_owned()).unwrap();
@@ -689,6 +1301,58 @@ mod tests {
);
}
#[test]
fn diversity_penalty_can_surface_a_distinct_result() {
assert_eq!(diversity_rank(0.94, REDUNDANCY_PENALTY_START), 0.94);
assert!(diversity_rank(0.94, 0.0) > diversity_rank(0.99, 1.0));
assert!(diversity_rank(0.99, 0.90) > diversity_rank(0.94, 0.0));
}
#[test]
fn semantic_similarity_can_rescue_a_visually_different_match() {
let semantic_match = hybrid_similarity_score(0.20, Some(0.10));
assert!(semantic_match.included);
assert!(semantic_match.relevance > MIN_HYBRID_SCORE);
let unrelated = hybrid_similarity_score(0.20, Some(0.04));
assert!(!unrelated.included);
let fallback = hybrid_similarity_score(0.70, None);
assert!(fallback.included);
assert_eq!(fallback.relevance, 0.70);
}
#[test]
fn natural_language_search_ranks_shared_text_image_embeddings() {
let record = |id: &str, name: &str| {
Arc::new(ImageRecord {
id: id.into(),
path: name.into(),
relative_path: name.into(),
search_key: name.to_lowercase(),
name: name.into(),
album: String::new(),
width: 10,
height: 10,
size: 1,
modified_ms: 1,
modified_ns: 1,
})
};
let query = Arc::new(TextSearchEmbedding::for_test([(0, 100)]));
let red = Arc::new(TextSearchEmbedding::for_test([(0, 100)]));
let blue = Arc::new(TextSearchEmbedding::for_test([(1, 100)]));
let ranked = rank_text_search(
query,
vec![
(record("blue", "first.png"), Some(blue), false),
(record("red", "second.png"), Some(red), false),
],
);
assert_eq!(ranked[0].id, "red");
assert_eq!(ranked[1].id, "blue");
}
#[test]
fn normalizes_image_queries_once() {
let query = ImagesQuery {
@@ -729,6 +1393,20 @@ mod tests {
);
}
#[tokio::test]
async fn status_exposes_optional_semantic_index_progress() {
let response = request("/api/status").await;
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let status: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(status["semantic"]["enabled"], false);
assert_eq!(status["semantic"]["backgroundComplete"], true);
assert_eq!(status["textSearch"]["enabled"], false);
assert_eq!(status["textSearch"]["backgroundComplete"], true);
}
#[tokio::test]
async fn unknown_api_and_image_routes_return_not_found() {
assert_eq!(
@@ -743,6 +1421,10 @@ mod tests {
request("/api/images/missing/original").await.status(),
StatusCode::NOT_FOUND
);
assert_eq!(
request("/api/images/missing/similar").await.status(),
StatusCode::NOT_FOUND
);
let response = request("/missing").await;
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(
@@ -753,6 +1435,196 @@ mod tests {
assert_eq!(response.headers()[header::REFERRER_POLICY], "no-referrer");
}
#[tokio::test]
async fn similar_images_are_ranked_exclude_the_source_and_paginate() {
let temp = tempfile::tempdir().unwrap();
let gallery = temp.path().join("gallery");
std::fs::create_dir(&gallery).unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([220, 80, 40]))
.save(gallery.join("source.png"))
.unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([216, 82, 43]))
.save(gallery.join("closest.png"))
.unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([205, 88, 48]))
.save(gallery.join("nearby.png"))
.unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([20, 50, 220]))
.save(gallery.join("blue.png"))
.unwrap();
let index = crate::gallery::scan_gallery(&gallery, None).unwrap();
let source_id = index
.images
.iter()
.find(|image| image.name == "source.png")
.unwrap()
.id
.clone();
let thumbnails = ThumbnailManager::new(temp.path().join("cache")).unwrap();
thumbnails.reconcile(&index.images);
thumbnails.start_workers(1);
for image in &index.images {
thumbnails.ensure_ready(&image.id).await.unwrap();
}
let app = router(AppState {
index: Arc::new(RwLock::new(Arc::new(index))),
thumbnails,
});
let response = app
.clone()
.oneshot(
Request::builder()
.uri(format!("/api/images/{source_id}/similar?offset=0&limit=1"))
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let page: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(page["total"], 2);
assert_eq!(page["nextOffset"], 1);
assert_eq!(page["items"].as_array().unwrap().len(), 1);
assert_eq!(page["items"][0]["name"], "closest.png");
assert_ne!(page["items"][0]["id"], source_id);
let response = app
.oneshot(
Request::builder()
.uri(format!("/api/images/{source_id}/similar?offset=1&limit=1"))
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let page: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(page["total"], 2);
assert_eq!(page["nextOffset"], serde_json::Value::Null);
assert_eq!(page["items"][0]["name"], "nearby.png");
}
#[tokio::test]
async fn similar_images_omit_unrelated_candidates() {
let temp = tempfile::tempdir().unwrap();
let gallery = temp.path().join("gallery");
std::fs::create_dir(&gallery).unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([220, 80, 40]))
.save(gallery.join("source.png"))
.unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([20, 50, 220]))
.save(gallery.join("unrelated.png"))
.unwrap();
let index = crate::gallery::scan_gallery(&gallery, None).unwrap();
let source_id = index
.images
.iter()
.find(|image| image.name == "source.png")
.unwrap()
.id
.clone();
let thumbnails = ThumbnailManager::new(temp.path().join("cache")).unwrap();
thumbnails.reconcile(&index.images);
thumbnails.start_workers(1);
for image in &index.images {
thumbnails.ensure_ready(&image.id).await.unwrap();
}
let app = router(AppState {
index: Arc::new(RwLock::new(Arc::new(index))),
thumbnails,
});
let response = app
.oneshot(
Request::builder()
.uri(format!("/api/images/{source_id}/similar"))
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let page: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(page["total"], 0);
assert_eq!(page["nextOffset"], serde_json::Value::Null);
assert!(page["items"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn natural_language_api_searches_image_content_not_only_filenames() {
let temp = tempfile::tempdir().unwrap();
let gallery = temp.path().join("gallery");
std::fs::create_dir(&gallery).unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([230, 25, 20]))
.save(gallery.join("first.png"))
.unwrap();
image::RgbImage::from_pixel(40, 30, image::Rgb([20, 30, 230]))
.save(gallery.join("second.png"))
.unwrap();
let index = crate::gallery::scan_gallery(&gallery, None).unwrap();
let text_search = crate::text_search::TextSearchIndex::with_test_model(
Box::new(ColourTextSearchModel),
[44; 16],
);
let thumbnails = ThumbnailManager::new_with_indexes(
temp.path().join("cache"),
None,
Some(Arc::clone(&text_search)),
)
.unwrap();
thumbnails.reconcile(&index.images);
thumbnails.start_workers(1);
text_search.start_worker();
for image in &index.images {
thumbnails.ensure_ready(&image.id).await.unwrap();
}
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let status = thumbnails.status();
if status.background_complete && status.text_search.background_complete {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
})
.await
.expect("text-search index timed out");
let app = router(AppState {
index: Arc::new(RwLock::new(Arc::new(index))),
thumbnails,
});
let response = app
.oneshot(
Request::builder()
.uri("/api/images?search=%E7%BA%A2%E8%89%B2%E5%9B%BE%E7%89%87&limit=2")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let page: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(page["items"][0]["name"], "first.png");
assert_eq!(page["items"][1]["name"], "second.png");
}
#[tokio::test]
async fn indexed_images_expose_original_with_immutable_cache() {
let temp = tempfile::tempdir().unwrap();