diff --git a/ui/package-lock.json b/ui/package-lock.json index 4b53eba3..50100f42 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -24,6 +24,7 @@ "react-global-hooks": "^1.3.5", "react-icons": "^5.5.0", "react-select": "^5.10.1", + "react-virtuoso": "^4.18.7", "react-zoom-pan-pinch": "^4.0.3", "sqlite3": "^5.1.7", "systeminformation": "^5.27.11", @@ -5167,6 +5168,16 @@ "react-dom": ">=16.6.0" } }, + "node_modules/react-virtuoso": { + "version": "4.18.7", + "resolved": "https://registry.npmjs.org/react-virtuoso/-/react-virtuoso-4.18.7.tgz", + "integrity": "sha512-xNF5zDGEEIMB7cKwcen/pLig0YDf6OnfFrVgKFa7sHPf9fRem0CaLshyObbBcP88jzn0enavL39EgplgdyT21g==", + "license": "MIT", + "peerDependencies": { + "react": ">=16 || >=17 || >= 18 || >= 19", + "react-dom": ">=16 || >=17 || >= 18 || >=19" + } + }, "node_modules/react-zoom-pan-pinch": { "version": "4.0.3", "resolved": "https://registry.npmjs.org/react-zoom-pan-pinch/-/react-zoom-pan-pinch-4.0.3.tgz", diff --git a/ui/package.json b/ui/package.json index af33b910..86d990a4 100644 --- a/ui/package.json +++ b/ui/package.json @@ -28,6 +28,7 @@ "react-global-hooks": "^1.3.5", "react-icons": "^5.5.0", "react-select": "^5.10.1", + "react-virtuoso": "^4.18.7", "react-zoom-pan-pinch": "^4.0.3", "sqlite3": "^5.1.7", "systeminformation": "^5.27.11", diff --git a/ui/src/app/api/audio/art/[...audioPath]/route.ts b/ui/src/app/api/audio/art/[...audioPath]/route.ts index de187f9f..fbc65047 100644 --- a/ui/src/app/api/audio/art/[...audioPath]/route.ts +++ b/ui/src/app/api/audio/art/[...audioPath]/route.ts @@ -94,9 +94,9 @@ function extractArtFromTag(buf: Buffer): ArtResult { verMajor === 4 ? synchsafeToInt(tagData[offset + 4], tagData[offset + 5], tagData[offset + 6], tagData[offset + 7]) : (tagData[offset + 4] << 24) | - (tagData[offset + 5] << 16) | - (tagData[offset + 6] << 8) | - tagData[offset + 7]; + (tagData[offset + 5] << 16) | + (tagData[offset + 6] << 8) | + tagData[offset + 7]; const flag2 = tagData[offset + 9]; offset += 10; if (!id.trim() || size <= 0 || offset + size > tagData.length) break; @@ -135,19 +135,22 @@ export async function GET(request: NextRequest, { params }: { params: { audioPat const trainingRoot = await getTrainingFolder(); const dataRoot = await getDataRoot(); const allowedDirs = [datasetRoot, trainingRoot, dataRoot]; - const isAllowed = allowedDirs.some(d => filepath.startsWith(d)) && !filepath.includes('..'); + // Resolve so `..` segments collapse, then verify still under an allowed root. + // Substring `.includes('..')` false-positives on filenames containing `..` as text. + const resolved = path.resolve(filepath); + const isAllowed = allowedDirs.some(d => resolved === d || resolved.startsWith(d + path.sep)); if (!isAllowed) { return new NextResponse('Access denied', { status: 403 }); } - const stat = await fs.promises.stat(filepath).catch(() => null); + const stat = await fs.promises.stat(resolved).catch(() => null); if (!stat || !stat.isFile()) { return new NextResponse('File not found', { status: 404 }); } // Read only the ID3 tag (first min(tagSize, 4MB) bytes). // First read 10 bytes to get tag size, then read the full tag. - const fd = await fs.promises.open(filepath, 'r'); + const fd = await fs.promises.open(resolved, 'r'); try { const headerBuf = Buffer.alloc(10); await fd.read(headerBuf, 0, 10, 0); @@ -167,7 +170,7 @@ export async function GET(request: NextRequest, { params }: { params: { audioPat return new NextResponse('No album art found', { status: 404 }); } - return new NextResponse(art.data, { + return new NextResponse(art.data as any, { headers: { 'Content-Type': art.mime, 'Content-Length': String(art.data.length), diff --git a/ui/src/app/api/caption/get/route.ts b/ui/src/app/api/caption/get/route.ts index 968624ec..e1b202ba 100644 --- a/ui/src/app/api/caption/get/route.ts +++ b/ui/src/app/api/caption/get/route.ts @@ -4,6 +4,11 @@ import fs from 'fs'; import path from 'path'; import { getDatasetsRoot } from '@/server/settings'; +function isUnderRoot(filepath: string, root: string): boolean { + const resolved = path.resolve(filepath); + return resolved === root || resolved.startsWith(root + path.sep); +} + export async function POST(request: NextRequest) { let body; try { @@ -30,8 +35,10 @@ export async function POST(request: NextRequest) { // Get allowed directories const allowedDir = await getDatasetsRoot(); - // Security check: Ensure path is in allowed directory - const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..'); + // Security check: resolve so `..` segments collapse, then verify it's still + // under the allowed root. Substring `.includes('..')` would false-positive + // on filenames that contain `..` as text (e.g. an ellipsis in a filename). + const isAllowed = isUnderRoot(filepath, allowedDir); if (!isAllowed) { console.warn(`Access denied: ${filepath} not in ${allowedDir}`); diff --git a/ui/src/app/api/caption/getBatch/route.ts b/ui/src/app/api/caption/getBatch/route.ts new file mode 100644 index 00000000..281c16af --- /dev/null +++ b/ui/src/app/api/caption/getBatch/route.ts @@ -0,0 +1,45 @@ +/* eslint-disable */ +import { NextRequest, NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot } from '@/server/settings'; + +function isUnderRoot(filepath: string, root: string): boolean { + const resolved = path.resolve(filepath); + return resolved === root || resolved.startsWith(root + path.sep); +} + +export async function POST(request: NextRequest) { + let body; + try { + body = await request.json(); + } catch { + return new NextResponse(null, { status: 499 }); + } + + if (request.signal.aborted) { + return new NextResponse(null, { status: 499 }); + } + + const { imgPaths } = body as { imgPaths?: string[] }; + if (!Array.isArray(imgPaths)) { + return NextResponse.json({ error: 'imgPaths must be an array' }, { status: 400 }); + } + + const allowedDir = await getDatasetsRoot(); + const captions: Record = {}; + + for (const imgPath of imgPaths) { + if (typeof imgPath !== 'string') continue; + if (!isUnderRoot(imgPath, allowedDir)) continue; + + const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt'; + try { + captions[imgPath] = fs.existsSync(captionPath) ? fs.readFileSync(captionPath, 'utf-8') : ''; + } catch { + captions[imgPath] = ''; + } + } + + return NextResponse.json({ captions }); +} diff --git a/ui/src/app/api/datasets/listImages/route.ts b/ui/src/app/api/datasets/listImages/route.ts index 7f144f8c..92e666a8 100644 --- a/ui/src/app/api/datasets/listImages/route.ts +++ b/ui/src/app/api/datasets/listImages/route.ts @@ -18,6 +18,9 @@ export async function POST(request: Request) { // Find all images recursively const imageFiles = findImagesRecursively(datasetFolder); + // Sort server-side so the client doesn't have to sort large lists + imageFiles.sort((a, b) => a.localeCompare(b)); + // Format response const result = imageFiles.map(imgPath => ({ img_path: imgPath, @@ -39,19 +42,20 @@ function findImagesRecursively(dir: string): string[] { const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp', '.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv', '.mp3', '.wav', '.flac', '.ogg']; let results: string[] = []; - const items = fs.readdirSync(dir); + // withFileTypes avoids a separate statSync per entry — a big win on large datasets + const entries = fs.readdirSync(dir, { withFileTypes: true }); - for (const item of items) { - const itemPath = path.join(dir, item); - const stat = fs.statSync(itemPath); + for (const entry of entries) { + const name = entry.name; + if (name.startsWith('.')) continue; + const itemPath = path.join(dir, name); - if (stat.isDirectory() && item !== '_controls' && !item.startsWith('.')) { - // If it's a directory, recursively search it + if (entry.isDirectory()) { + if (name === '_controls') continue; results = results.concat(findImagesRecursively(itemPath)); - } else { - // If it's a file, check if it's an image - const ext = path.extname(itemPath).toLowerCase(); - if (imageExtensions.includes(ext) && !item.startsWith('.')) { + } else if (entry.isFile()) { + const ext = path.extname(name).toLowerCase(); + if (imageExtensions.includes(ext)) { results.push(itemPath); } } diff --git a/ui/src/app/api/files/[...filePath]/route.ts b/ui/src/app/api/files/[...filePath]/route.ts index 4d9f44b7..3ecaf66f 100644 --- a/ui/src/app/api/files/[...filePath]/route.ts +++ b/ui/src/app/api/files/[...filePath]/route.ts @@ -15,32 +15,36 @@ export async function GET(request: NextRequest, { params }: { params: { filePath const trainingRoot = await getTrainingFolder(); const allowedDirs = [datasetRoot, trainingRoot]; - // Security check: Ensure path is in allowed directory - const isAllowed = - allowedDirs.some(allowedDir => decodedFilePath.startsWith(allowedDir)) && !decodedFilePath.includes('..'); + // Security check: resolve so `..` segments collapse, then verify still under + // an allowed root. Substring `.includes('..')` false-positives on filenames + // containing `..` as text (e.g. an ellipsis in a filename). + const resolvedFilePath = path.resolve(decodedFilePath); + const isAllowed = allowedDirs.some( + allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep), + ); if (!isAllowed) { - console.warn(`Access denied: ${decodedFilePath} not in ${allowedDirs.join(', ')}`); + console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`); return new NextResponse('Access denied', { status: 403 }); } // Check if file exists - if (!fs.existsSync(decodedFilePath)) { - console.warn(`File not found: ${decodedFilePath}`); + if (!fs.existsSync(resolvedFilePath)) { + console.warn(`File not found: ${resolvedFilePath}`); return new NextResponse('File not found', { status: 404 }); } // Get file info - const stat = fs.statSync(decodedFilePath); + const stat = fs.statSync(resolvedFilePath); if (!stat.isFile()) { return new NextResponse('Not a file', { status: 400 }); } // Get filename for Content-Disposition - const filename = path.basename(decodedFilePath); + const filename = path.basename(resolvedFilePath); // Determine content type - const ext = path.extname(decodedFilePath).toLowerCase(); + const ext = path.extname(resolvedFilePath).toLowerCase(); const contentTypeMap: { [key: string]: string } = { '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', @@ -87,7 +91,7 @@ export async function GET(request: NextRequest, { params }: { params: { filePath const end = parts[1] ? parseInt(parts[1], 10) : Math.min(start + 10 * 1024 * 1024, stat.size - 1); // 10MB chunks const chunkSize = end - start + 1; - const fileStream = fs.createReadStream(decodedFilePath, { + const fileStream = fs.createReadStream(resolvedFilePath, { start, end, highWaterMark: 64 * 1024, // 64KB buffer @@ -103,7 +107,7 @@ export async function GET(request: NextRequest, { params }: { params: { filePath }); } else { // For full file download, read directly without streaming wrapper - const fileStream = fs.createReadStream(decodedFilePath, { + const fileStream = fs.createReadStream(resolvedFilePath, { highWaterMark: 64 * 1024, // 64KB buffer }); diff --git a/ui/src/app/api/img/[...imagePath]/route.ts b/ui/src/app/api/img/[...imagePath]/route.ts index 8bbf358d..a3877e8a 100644 --- a/ui/src/app/api/img/[...imagePath]/route.ts +++ b/ui/src/app/api/img/[...imagePath]/route.ts @@ -41,21 +41,26 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat const allowedDirs = [datasetRoot, trainingRoot, dataRoot]; - // Security check: Ensure path is in allowed directory - const isAllowed = allowedDirs.some(allowedDir => filepath.startsWith(allowedDir)) && !filepath.includes('..'); + // Security check: resolve the path so any `..` segments are collapsed, + // then ensure it's still under an allowed root. (Plain `.includes('..')` + // false-positives on filenames that contain `..` as text, e.g. an ellipsis.) + const resolved = path.resolve(filepath); + const isAllowed = allowedDirs.some( + allowedDir => resolved === allowedDir || resolved.startsWith(allowedDir + path.sep), + ); if (!isAllowed) { - console.warn(`Access denied: ${filepath} not in ${allowedDirs.join(', ')}`); + console.warn(`Access denied: ${resolved} not in ${allowedDirs.join(', ')}`); return new NextResponse('Access denied', { status: 403 }); } // Stat file (async) - const stat = await fs.promises.stat(filepath).catch(() => null); + const stat = await fs.promises.stat(resolved).catch(() => null); if (!stat || !stat.isFile()) { return new NextResponse('File not found', { status: 404 }); } - const ext = path.extname(filepath).toLowerCase(); + const ext = path.extname(resolved).toLowerCase(); const contentType = contentTypeMap[ext] || 'application/octet-stream'; // Support range requests for video/audio seeking @@ -66,7 +71,7 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat const end = parts[1] ? parseInt(parts[1], 10) : stat.size - 1; const chunkSize = end - start + 1; - const stream = fs.createReadStream(filepath, { start, end }); + const stream = fs.createReadStream(resolved, { start, end }); const readable = new ReadableStream({ start(controller) { stream.on('data', chunk => controller.enqueue(chunk)); @@ -91,7 +96,7 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat } // Stream the file instead of buffering it entirely - const stream = fs.createReadStream(filepath); + const stream = fs.createReadStream(resolved); const readable = new ReadableStream({ start(controller) { stream.on('data', chunk => controller.enqueue(chunk)); diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx index ebcb47ee..e0b03de3 100644 --- a/ui/src/app/datasets/[datasetName]/page.tsx +++ b/ui/src/app/datasets/[datasetName]/page.tsx @@ -1,8 +1,9 @@ 'use client'; -import { useEffect, useState, use, useMemo } from 'react'; +import { useEffect, useState, use, useMemo, useCallback } from 'react'; import { LuImageOff, LuLoader, LuBan } from 'react-icons/lu'; import { FaChevronLeft } from 'react-icons/fa'; +import { VirtuosoGrid } from 'react-virtuoso'; import DatasetImageCard from '@/components/DatasetImageCard'; import DatasetImageViewer from '@/components/DatasetImageViewer'; import { Button } from '@headlessui/react'; @@ -22,17 +23,16 @@ export default function DatasetPage({ params }: { params: { datasetName: string const { settings, isSettingsLoaded } = useSettings(); const [selectedImgPath, setSelectedImgPath] = useState(null); const [captionRefreshKeys, setCaptionRefreshKeys] = useState>({}); + const [scrollParent, setScrollParent] = useState(null); + const scrollParentCallback = useCallback((el: HTMLDivElement | null) => setScrollParent(el), []); const refreshImageList = (dbName: string) => { setStatus('loading'); - console.log('Fetching images for dataset:', dbName); apiClient .post('/api/datasets/listImages', { datasetName: dbName }) .then((res: any) => { const data = res.data; - console.log('Images:', data.images); - // sort - data.images.sort((a: { img_path: string }, b: { img_path: string }) => a.img_path.localeCompare(b.img_path)); + // Server already sorts; avoid the client-side sort that's expensive on large lists. setImgList(data.images); setStatus('success'); }) @@ -43,6 +43,8 @@ export default function DatasetPage({ params }: { params: { datasetName: string }; useOpenImagesModalOnDrag(datasetName, () => refreshImageList(datasetName)); + const imgPaths = useMemo(() => imgList.map(img => img.img_path), [imgList]); + useEffect(() => { if (datasetName) { refreshImageList(datasetName); @@ -129,28 +131,36 @@ export default function DatasetPage({ params }: { params: { datasetName: string - + {PageInfoContent} - {status === 'success' && imgList.length > 0 && ( -
- {imgList.map(img => ( - refreshImageList(datasetName)} - onImageClick={() => setSelectedImgPath(img.img_path)} - captionRefreshKey={captionRefreshKeys[img.img_path] || 0} - /> - ))} -
+ {status === 'success' && imgList.length > 0 && scrollParent && ( + { + const img = imgList[index]; + if (!img) return null; + return ( + refreshImageList(datasetName)} + onImageClick={() => setSelectedImgPath(img.img_path)} + captionRefreshKey={captionRefreshKeys[img.img_path] || 0} + /> + ); + }} + computeItemKey={index => imgList[index]?.img_path ?? index} + /> )} img.img_path)} + imageList={imgPaths} onChange={setSelectedImgPath} refreshImages={() => refreshImageList(datasetName)} onCaptionSaved={path => setCaptionRefreshKeys(prev => ({ ...prev, [path]: (prev[path] || 0) + 1 }))} diff --git a/ui/src/components/DatasetImageCard.tsx b/ui/src/components/DatasetImageCard.tsx index 8a8112cd..bedb42b2 100644 --- a/ui/src/components/DatasetImageCard.tsx +++ b/ui/src/components/DatasetImageCard.tsx @@ -1,10 +1,11 @@ -import React, { useRef, useEffect, useState, ReactNode, KeyboardEvent } from 'react'; -import { FaTrashAlt, FaEye, FaEyeSlash } from 'react-icons/fa'; +import React, { useEffect, useState, ReactNode, KeyboardEvent, useRef } from 'react'; +import { FaTrashAlt } from 'react-icons/fa'; import { openConfirm } from './ConfirmModal'; import classNames from 'classnames'; import { apiClient } from '@/utils/api'; import AudioPlayer from './AudioPlayer'; import { isVideo, isAudio } from '@/utils/basic'; +import useCaptionBatch, { setCachedCaption } from '@/hooks/useCaptionBatch'; interface DatasetImageCardProps { imageUrl: string; @@ -27,131 +28,85 @@ const DatasetImageCard: React.FC = ({ onImageClick, captionRefreshKey = 0, }) => { - const cardRef = useRef(null); - const [isVisible, setIsVisible] = useState(false); - const [inViewport, setInViewport] = useState(false); const [loaded, setLoaded] = useState(false); - const [isCaptionLoaded, setIsCaptionLoaded] = useState(false); + const [showAudioPlayer, setShowAudioPlayer] = useState(true); + const [pollTick, setPollTick] = useState(0); + + const combinedRefreshKey = captionRefreshKey + pollTick; + const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(imageUrl, combinedRefreshKey); + const [caption, setCaption] = useState(''); const [savedCaption, setSavedCaption] = useState(''); - const abortControllerRef = useRef(null); + const dirtyRef = useRef(false); - const fetchCaption = async () => { - if (isCaptionLoaded) return; - abortControllerRef.current?.abort(); - const controller = new AbortController(); - abortControllerRef.current = controller; - apiClient - .post(`/api/caption/get`, { imgPath: imageUrl }, { signal: controller.signal }) - .then(res => res.data) - .then(data => { - console.log('Caption fetched:', data); - if (data) { - data = `${data}`; - } - setCaption(data || ''); - setSavedCaption((data || '').trim()); - setIsCaptionLoaded(true); - }) - .catch(error => { - if (controller.signal.aborted) return; - console.error('Error fetching caption:', error); - }) - .finally(() => { - if (abortControllerRef.current === controller) { - abortControllerRef.current = null; - } - }); - }; + // Sync from the fetched caption, but don't clobber unsaved local edits. + useEffect(() => { + if (!isCaptionLoaded) return; + if (dirtyRef.current) return; + setCaption(fetchedCaption); + setSavedCaption(fetchedCaption.trim()); + }, [fetchedCaption, isCaptionLoaded]); + + // Poll while auto-captioning so backend-written captions show up. + useEffect(() => { + if (!isAutoCaptioning) return; + const interval = setInterval(() => setPollTick(t => t + 1), 5000); + return () => clearInterval(interval); + }, [isAutoCaptioning]); const saveCaption = () => { const trimmedCaption = caption.trim(); - if (trimmedCaption === savedCaption) return; + if (trimmedCaption === savedCaption) { + dirtyRef.current = false; + return; + } apiClient .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption }) - .then(res => res.data) - .then(data => { - console.log('Caption saved:', data); + .then(() => { setSavedCaption(trimmedCaption); + setCachedCaption(imageUrl, trimmedCaption); + dirtyRef.current = false; }) .catch(error => { console.error('Error saving caption:', error); }); }; - // Only fetch caption when the component is both in viewport and visible + // Save any pending edit if the card unmounts (e.g. scrolled out of the virtualized window). + const latestRef = useRef({ caption, savedCaption, imageUrl }); useEffect(() => { - if (inViewport && isVisible) { - fetchCaption(); - } - }, [inViewport, isVisible, isCaptionLoaded]); - - // Poll for caption updates every 5 seconds while auto-captioning + latestRef.current = { caption, savedCaption, imageUrl }; + }); useEffect(() => { - if (!isAutoCaptioning || !inViewport || !isVisible) return; - const interval = setInterval(() => { - // Reset so fetchCaption will re-fetch - setIsCaptionLoaded(false); - }, 5000); - return () => clearInterval(interval); - }, [isAutoCaptioning, inViewport, isVisible]); - - // External trigger (e.g. caption edited in the full-screen viewer) — re-fetch - useEffect(() => { - if (captionRefreshKey === 0) return; - setIsCaptionLoaded(false); - }, [captionRefreshKey]); - - useEffect(() => { - // Create intersection observer to check viewport visibility - const observer = new IntersectionObserver( - entries => { - if (entries[0].isIntersecting) { - setInViewport(true); - // Initialize isVisible to true when first coming into view - if (!isVisible) { - setIsVisible(true); - } - } else { - setInViewport(false); - // Cancel any in-flight caption fetch when scrolling away - abortControllerRef.current?.abort(); - } - }, - { threshold: 0.1 }, - ); - - if (cardRef.current) { - observer.observe(cardRef.current); - } - return () => { - observer.disconnect(); + if (!dirtyRef.current) return; + const { caption: c, savedCaption: s, imageUrl: url } = latestRef.current; + const trimmed = c.trim(); + if (trimmed === s) return; + apiClient + .post('/api/img/caption', { imgPath: url, caption: trimmed }) + .then(() => setCachedCaption(url, trimmed)) + .catch(err => console.error('Error saving caption on unmount:', err)); }; }, []); - const toggleVisibility = (): void => { - setIsVisible(prev => !prev); - if (!isVisible && !isCaptionLoaded) { - fetchCaption(); - } - }; - const handleLoad = (): void => { setLoaded(true); }; const handleKeyDown = (e: KeyboardEvent): void => { - // If Enter is pressed without Shift, prevent default behavior and save if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); saveCaption(); } }; - const isCaptionCurrent = caption.trim() === savedCaption; + const handleCaptionChange = (value: string) => { + dirtyRef.current = value.trim() !== savedCaption; + setCaption(value); + }; - const [showAudioPlayer, setShowAudioPlayer] = useState(true); + const isCaptionCurrent = caption.trim() === savedCaption; const isItAVideo = isVideo(imageUrl); const isItAudio = isAudio(imageUrl); @@ -159,66 +114,49 @@ const DatasetImageCard: React.FC = ({ return (
- {/* Square image container */} -
+
- {inViewport && isVisible && ( - <> - {isItAVideo && ( -