From aded9b5f89ebbfe32bf399ca750123bd4fdb6da1 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 6 Jul 2026 21:41:56 -0400 Subject: [PATCH 1/3] Add desktop video search / IQR backend Per-dataset search index lifecycle (build via a process_video.py job, status from an index_meta.json sidecar + ITQ file check, delete) and a persistent QueryServiceManager wrapping viame.core.query_service over NDJSON stdio (open index, formulate from image chip + boxes, query with optional warm-start model, refine with +/- feedback, export the SVM as a runnable trained pipeline in DIVE_Pipelines). One index open at a time; the embedded postgres is stopped gracefully on app cleanup. Co-Authored-By: Claude Fable 5 --- client/platform/desktop/backend/ipcService.ts | 68 ++ .../desktop/backend/native/videoSearch.ts | 602 ++++++++++++++++++ client/platform/desktop/background.ts | 4 + client/platform/desktop/constants.ts | 28 +- 4 files changed, 700 insertions(+), 2 deletions(-) create mode 100644 client/platform/desktop/backend/native/videoSearch.ts diff --git a/client/platform/desktop/backend/ipcService.ts b/client/platform/desktop/backend/ipcService.ts index 91a5664d1..c50cba838 100644 --- a/client/platform/desktop/backend/ipcService.ts +++ b/client/platform/desktop/backend/ipcService.ts @@ -14,6 +14,7 @@ import { ExportTrainedPipeline, ConversionArgs, DesktopJob, + BuildSearchIndex, } from 'platform/desktop/constants'; import { convertMedia } from 'platform/desktop/backend/native/mediaJobs'; import { closeChildById } from 'platform/desktop/backend/native/processManager'; @@ -28,6 +29,7 @@ import { listen } from './server'; import { getInteractiveServiceManager, } from './native/interactive'; +import * as videoSearch from './native/videoSearch'; import { SegmentationPredictRequest, SegmentationStereoSegmentRequest, @@ -459,4 +461,70 @@ export default function register() { const stereoService = getInteractiveServiceManager(); return { enabled: stereoService.isEnabled() }; }); + + /** + * Video Search / IQR Service + */ + + ipcMain.handle('video-search-installed', async () => ( + videoSearch.isVideoSearchInstalled(settings.get()) + )); + + ipcMain.handle('video-search-index-status', async (_, datasetId: string) => ( + videoSearch.getIndexStatus(settings.get(), datasetId) + )); + + ipcMain.handle('video-search-build-index', async (event, args: BuildSearchIndex) => { + const updater = (update: DesktopJobUpdate) => { + event.sender.send('job-update', update); + }; + // A rebuild invalidates whatever the service currently has open. + const manager = videoSearch.getQueryServiceManager(); + if (manager.openIndexDir() === videoSearch.getIndexDir(settings.get(), args.datasetId)) { + await manager.closeIndex(); + } + return videoSearch.buildIndex(settings.get(), args, updater); + }); + + ipcMain.handle('video-search-delete-index', async (_, datasetId: string) => { + await videoSearch.deleteIndex(settings.get(), datasetId); + return { success: true }; + }); + + ipcMain.handle('video-search-open-index', async (_, datasetId: string) => { + const currentSettings = settings.get(); + const status = await videoSearch.getIndexStatus(currentSettings, datasetId); + if (!status.built) { + throw new Error('No built search index exists for this dataset'); + } + const manager = videoSearch.getQueryServiceManager(); + await manager.openIndex(currentSettings, videoSearch.getIndexDir(currentSettings, datasetId)); + return { success: true }; + }); + + ipcMain.handle('video-search-formulate', async (_, args: { imagePath: string; boxes?: number[][] }) => { + const manager = videoSearch.getQueryServiceManager(); + return manager.formulateQuery(args.imagePath, args.boxes); + }); + + ipcMain.handle('video-search-query', async (_, args: { threshold?: number; iqrModelB64?: string }) => { + const manager = videoSearch.getQueryServiceManager(); + return manager.processQuery(args.threshold, args.iqrModelB64); + }); + + ipcMain.handle('video-search-refine', async (_, args: { positiveIds: number[]; negativeIds: number[] }) => { + const manager = videoSearch.getQueryServiceManager(); + return manager.refine(args.positiveIds, args.negativeIds); + }); + + ipcMain.handle('video-search-export-model', async (_, args: { name: string }) => { + const outputDir = await videoSearch.exportSearchModel(settings.get(), args.name); + return { success: true, outputDir }; + }); + + ipcMain.handle('video-search-close', async () => { + const manager = videoSearch.getQueryServiceManager(); + await manager.closeIndex(); + return { success: true }; + }); } diff --git a/client/platform/desktop/backend/native/videoSearch.ts b/client/platform/desktop/backend/native/videoSearch.ts new file mode 100644 index 000000000..94f97ab10 --- /dev/null +++ b/client/platform/desktop/backend/native/videoSearch.ts @@ -0,0 +1,602 @@ +/** + * Video Search / IQR backend for Desktop + * + * Provides per-dataset search index management (build via a process_video.py + * job, status, delete) and a persistent query service manager + * (viame.core.query_service) that holds the KWIVER query/IQR pipeline open so + * refinement iterations are interactive. One index may be open at a time (the + * embedded PostgreSQL instance binds a fixed port), enforced here by closing + * the previous index before opening another. + * + * On-disk layout (per dataset): + * /DIVE_Projects//search_index/ + * index_meta.json - written by DIVE on successful build + * ingest_list.txt - media list handed to process_video.py + * database/ - embedded PostgreSQL data + ITQ/LSH index files + */ + +import OS from 'os'; +import { spawn, ChildProcess } from 'child_process'; +import npath from 'path'; +import readline from 'readline'; +import fs from 'fs-extra'; +import { EventEmitter } from 'events'; + +import { + Settings, DesktopJob, DesktopJobUpdater, + SearchIndexMeta, BuildSearchIndex, +} from 'platform/desktop/constants'; +import { serialize } from 'platform/desktop/backend/serializers/viame'; +import { observeChild } from './processManager'; +import * as common from './common'; +import { jobFileEchoMiddleware, createCustomWorkingDirectory } from './utils'; +import linux from './linux'; +import win32 from './windows'; + +const SearchIndexFolderName = 'search_index'; +const IndexMetaFileName = 'index_meta.json'; +const QueryPipelineName = 'query_retrieval_and_iqr.pipe'; +const ExportTemplateName = npath.join('templates', 'detector_generic_svm.pipe'); + +const IndexPipelines: Record = { + detections: 'index_default.pipe', + tracking: 'index_default.trk.pipe', + existing: 'index_existing.pipe', +}; + +function getCurrentPlatform() { + return OS.platform() === 'win32' ? win32 : linux; +} + +/** + * Video search requires the query pipeline (built with database support) and + * the bundled PostgreSQL binaries used to host the descriptor index. + */ +async function isVideoSearchInstalled(settings: Settings): Promise { + const pipelines = npath.join(settings.viamePath, 'configs', 'pipelines'); + const initdb = OS.platform() === 'win32' ? 'initdb.exe' : 'initdb'; + const checks = await Promise.all([ + fs.pathExists(npath.join(pipelines, QueryPipelineName)), + fs.pathExists(npath.join(pipelines, 'sql_init_table.sql')), + fs.pathExists(npath.join(settings.viamePath, 'bin', initdb)), + ]); + return checks.every((c) => c); +} + +function getIndexDir(settings: Settings, datasetId: string): string { + return npath.join(settings.dataPath, 'DIVE_Projects', datasetId, SearchIndexFolderName); +} + +export interface SearchIndexStatus { + exists: boolean; + built: boolean; + meta?: SearchIndexMeta; +} + +async function getIndexStatus(settings: Settings, datasetId: string): Promise { + const indexDir = getIndexDir(settings, datasetId); + const metaPath = npath.join(indexDir, IndexMetaFileName); + if (!(await fs.pathExists(metaPath))) { + return { exists: await fs.pathExists(indexDir), built: false }; + } + const meta = (await fs.readJson(metaPath)) as SearchIndexMeta; + // Sanity: the ITQ model files must exist for the index to be queryable + const itqDir = npath.join(indexDir, 'database', 'ITQ'); + const built = (await fs.pathExists(itqDir)) + && (await fs.readdir(itqDir)).some((f) => f.startsWith('itq.model')); + return { exists: true, built, meta }; +} + +/** + * Build (or rebuild) the search index for a dataset by running + * process_video.py --init --build-index as a DIVE job. + */ +async function buildIndex( + settings: Settings, + args: BuildSearchIndex, + updater: DesktopJobUpdater, +): Promise { + const { datasetId, method } = args; + const platform = getCurrentPlatform(); + const isValid = await platform.validateViamePath(settings); + if (isValid !== true) { + throw new Error(isValid); + } + + const projectInfo = await common.getValidatedProjectDir(settings, datasetId); + const meta = await common.loadJsonMetadata(projectInfo.metaFileAbsPath); + if (meta.multiCam) { + throw new Error('Search indexes are not yet supported on multi-camera datasets'); + } + + const indexDir = getIndexDir(settings, datasetId); + // Rebuild from scratch: the DB re-init inside process_video handles the + // database folder, but a stale meta file must not mark a failed build as + // usable. + await fs.remove(npath.join(indexDir, IndexMetaFileName)); + await fs.ensureDir(indexDir); + + // Media list: image sequences list every frame; videos list the one file. + const ingestList = npath.join(indexDir, 'ingest_list.txt'); + if (meta.type === 'video') { + const videoAbsPath = npath.join(meta.originalBasePath, meta.originalVideoFile); + await fs.writeFile(ingestList, `${videoAbsPath}\n`); + } else { + const fileData = meta.originalImageFiles + .map((f) => npath.join(meta.originalBasePath, f)) + .join('\n'); + await fs.writeFile(ingestList, `${fileData}\n`); + } + + const viameConstants = platform.getViameConstants(settings); + const pythonExe = platform.getViamePythonExe(settings); + const processVideo = npath.join(settings.viamePath, 'configs', 'process_video.py'); + + const command = [ + `${viameConstants.setupScriptAbs} &&`, + `"${pythonExe}" "${processVideo}"`, + '--init --no-reset-prompt', + `-l "${ingestList}"`, + `-p "pipelines/${IndexPipelines[method]}"`, + '-o database', + '--build-index', + `-install "${settings.viamePath}"`, + ]; + if (meta.type === 'video') { + command.push(`-frate ${meta.fps}`); + } + + // Index around this dataset's existing annotations + if (method === 'existing') { + const detectionsCsv = npath.join(indexDir, 'input_detections.csv'); + const csvStream = fs.createWriteStream(detectionsCsv); + const inputData = await common.loadAnnotationFile(projectInfo.trackFileAbsPath); + await serialize(csvStream, inputData, meta); + csvStream.end(); + command.push(`-id "${detectionsCsv}"`); + } + + const jobWorkDir = await createCustomWorkingDirectory(settings, 'search_index', datasetId.replace('/', '_')); + const joblog = npath.join(jobWorkDir, 'runlog.txt'); + + const job = observeChild(spawn(command.join(' '), { + shell: viameConstants.shell, + cwd: indexDir, + })); + + const jobBase: DesktopJob = { + key: `search_index_${job.pid}_${jobWorkDir}`, + command: command.join(' '), + jobType: 'pipeline', + pid: job.pid, + args, + title: `Build search index (${method})`, + workingDir: jobWorkDir, + datasetIds: [datasetId], + exitCode: job.exitCode, + startTime: new Date(), + }; + + fs.writeFile(npath.join(jobWorkDir, 'dive_job_manifest.json'), JSON.stringify(jobBase, null, 2)); + + updater({ ...jobBase, body: [''] }); + + job.stdout.on('data', jobFileEchoMiddleware(jobBase, updater, joblog)); + job.stderr.on('data', jobFileEchoMiddleware(jobBase, updater, joblog)); + + job.on('exit', async (code) => { + if (code === 0) { + const indexMeta: SearchIndexMeta = { + version: 1, + method, + fps: meta.fps, + datasets: [datasetId], + createdAt: (new Date()).toISOString(), + }; + try { + await fs.writeJson(npath.join(indexDir, IndexMetaFileName), indexMeta, { spaces: 2 }); + } catch (err) { + console.error('Failed to write search index metadata', err); + } + } + updater({ + ...jobBase, body: [''], exitCode: code, endTime: new Date(), + }); + }); + return jobBase; +} + +/** Loose shape of a JSON response line from the Python query service. */ +interface ServiceResponse { + id: string; + success?: boolean; + error?: string; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + [key: string]: any; +} + +interface PendingRequest { + resolve: (response: ServiceResponse) => void; + reject: (error: Error) => void; + timeout: NodeJS.Timeout; +} + +/** + * Manages ONE persistent viame.core.query_service subprocess speaking + * newline-delimited JSON. Modeled on InteractiveServiceManager. + */ +export class QueryServiceManager extends EventEmitter { + private process: ChildProcess | null = null; + + private readline: readline.Interface | null = null; + + private pendingRequests: Map = new Map(); + + private isStarting = false; + + private startPromise: Promise | null = null; + + private requestCounter = 0; + + /** Index directory currently opened in the service, if any. */ + private currentIndexDir: string | null = null; + + // First open loads the descriptor index + pipeline; queries train SVMs. + private readonly requestTimeoutMs = 300000; + + isReady(): boolean { + return this.process !== null && this.process.exitCode === null; + } + + openIndexDir(): string | null { + return this.isReady() ? this.currentIndexDir : null; + } + + private generateRequestId(): string { + this.requestCounter += 1; + return `req_${Date.now()}_${this.requestCounter}`; + } + + async ensureStarted(settings: Settings): Promise { + if (this.isReady()) { + return; + } + if (this.isStarting && this.startPromise) { + await this.startPromise; + return; + } + this.isStarting = true; + this.startPromise = this.doStart(settings); + try { + await this.startPromise; + } finally { + this.isStarting = false; + } + } + + private async doStart(settings: Settings): Promise { + await this.shutdown(); + + const platform = getCurrentPlatform(); + const isValid = await platform.validateViamePath(settings); + if (isValid !== true) { + throw new Error(isValid); + } + + return new Promise((resolve, reject) => { + const viameConstants = platform.getViameConstants(settings); + const pythonExe = platform.getViamePythonExe(settings); + const pyCommand = `"${pythonExe}" -s -m viame.core.query_service`; + const command = `${viameConstants.setupScriptAbs} && ${pyCommand}`; + + // eslint-disable-next-line no-console + console.log(`[VideoSearch] Starting query service: ${command}`); + + const stderrLines: string[] = []; + const maxStderrLines = 20; + + this.process = observeChild(spawn(command, { + shell: viameConstants.shell, + cwd: settings.viamePath, + stdio: ['pipe', 'pipe', 'pipe'], + })); + + if (this.process.stdout) { + this.readline = readline.createInterface({ + input: this.process.stdout, + crlfDelay: Infinity, + }); + this.readline.on('line', (line) => this.handleResponse(line)); + } + + let startupTimeout: NodeJS.Timeout | null = null; + + if (this.process.stderr) { + this.process.stderr.on('data', (data: Buffer) => { + const message = data.toString().trim(); + if (message) { + // eslint-disable-next-line no-console + console.log(`[VideoSearch] ${message}`); + stderrLines.push(message); + if (stderrLines.length > maxStderrLines) { + stderrLines.shift(); + } + if (message.includes('Service started, waiting for requests')) { + if (startupTimeout) { + clearTimeout(startupTimeout); + startupTimeout = null; + } + resolve(); + } + } + }); + } + + const rejectStartup = (reason: string) => { + if (!this.isStarting) { + return; + } + this.isStarting = false; + if (startupTimeout) { + clearTimeout(startupTimeout); + startupTimeout = null; + } + const details = stderrLines.filter((l) => l.trim()).slice(-8).join('\n '); + this.shutdown().finally(() => reject(new Error( + `Unable to start the video search service. ${reason}\n` + + `VIAME path: ${settings.viamePath}\n ${details}`, + ))); + }; + + this.process.on('exit', (code, signal) => { + // eslint-disable-next-line no-console + console.log(`[VideoSearch] Process exited with code ${code}, signal ${signal}`); + this.cleanup(); + if (this.isStarting) { + rejectStartup(`The service process exited before it became ready (code ${code}, signal ${signal}).`); + } + }); + + this.process.on('error', (err) => { + console.error('[VideoSearch] Process error:', err); + this.cleanup(); + if (this.isStarting) { + rejectStartup(`Failed to start the service process: ${err.message}`); + } + }); + + startupTimeout = setTimeout(() => { + rejectStartup('The service did not become ready within 5 minutes.'); + }, 300000); + }); + } + + private handleResponse(line: string): void { + let response: ServiceResponse; + try { + response = JSON.parse(line) as ServiceResponse; + } catch (err) { + console.error('[VideoSearch] Failed to parse response:', line, err); + return; + } + const pending = this.pendingRequests.get(response.id); + if (pending) { + clearTimeout(pending.timeout); + this.pendingRequests.delete(response.id); + pending.resolve(response); + } else { + console.warn(`[VideoSearch] Received response for unknown request: ${response.id}`); + } + } + + private sendRequest(payload: Record, timeoutLabel: string): Promise { + if (!this.isReady() || !this.process?.stdin) { + return Promise.reject(new Error('Video search service is not running')); + } + const id = this.generateRequestId(); + const fullRequest = { ...payload, id }; + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + this.pendingRequests.delete(id); + reject(new Error(`${timeoutLabel} request timed out after ${this.requestTimeoutMs}ms`)); + }, this.requestTimeoutMs); + + this.pendingRequests.set(id, { resolve, reject, timeout }); + + this.process!.stdin!.write(`${JSON.stringify(fullRequest)}\n`, (err) => { + if (err) { + clearTimeout(timeout); + this.pendingRequests.delete(id); + reject(err); + } + }); + }); + } + + private static check(response: ServiceResponse, label: string): ServiceResponse { + if (!response.success) { + throw new Error(response.error || `${label} failed`); + } + return response; + } + + // -------------------------------------------------------------- commands + + /** Open an index, closing any previously open one (one postgres at a time). */ + async openIndex(settings: Settings, indexDir: string): Promise { + await this.ensureStarted(settings); + if (this.currentIndexDir === indexDir) { + return; + } + if (this.currentIndexDir) { + await this.closeIndex(); + } + const response = await this.sendRequest({ command: 'open_index', index_dir: indexDir }, 'Open index'); + QueryServiceManager.check(response, 'Opening the search index'); + this.currentIndexDir = indexDir; + } + + async formulateQuery(imagePath: string, boxes?: number[][]): Promise { + const response = await this.sendRequest({ + command: 'formulate_query', + image_path: imagePath, + boxes, + }, 'Query formulation'); + return QueryServiceManager.check(response, 'Query formulation'); + } + + async processQuery(threshold?: number, iqrModelB64?: string): Promise { + const response = await this.sendRequest({ + command: 'process_query', + threshold: threshold ?? 0.0, + iqr_model_b64: iqrModelB64, + }, 'Query'); + return QueryServiceManager.check(response, 'Query'); + } + + async refine(positiveIds: number[], negativeIds: number[]): Promise { + const response = await this.sendRequest({ + command: 'refine', + positive_ids: positiveIds, + negative_ids: negativeIds, + }, 'Query refinement'); + return QueryServiceManager.check(response, 'Query refinement'); + } + + async exportModel(outputPath?: string): Promise { + const response = await this.sendRequest({ + command: 'export_model', + output_path: outputPath, + }, 'Model export'); + return QueryServiceManager.check(response, 'Model export'); + } + + async closeIndex(): Promise { + if (!this.isReady()) { + this.currentIndexDir = null; + return; + } + try { + await this.sendRequest({ command: 'close_index' }, 'Close index'); + } catch { + // best effort - shutting the process down also stops postgres + } + this.currentIndexDir = null; + } + + // ------------------------------------------------------------- lifecycle + + private cleanup(): void { + this.pendingRequests.forEach((pending) => { + clearTimeout(pending.timeout); + pending.reject(new Error('Video search service terminated')); + }); + this.pendingRequests.clear(); + + if (this.readline) { + this.readline.close(); + this.readline = null; + } + + this.process = null; + this.currentIndexDir = null; + this.emit('shutdown'); + } + + async shutdown(): Promise { + if (!this.process) { + return; + } + // eslint-disable-next-line no-console + console.log('[VideoSearch] Shutting down query service...'); + await new Promise((resolve) => { + const request = { id: this.generateRequestId(), command: 'shutdown' }; + if (this.process?.stdin?.writable) { + this.process.stdin.write(`${JSON.stringify(request)}\n`); + } + const timeoutId = setTimeout(() => { + if (this.process) { + // eslint-disable-next-line no-console + console.log('[VideoSearch] Force killing query service...'); + this.process.kill('SIGTERM'); + } + this.cleanup(); + resolve(); + }, 10000); + if (this.process) { + this.process.once('exit', () => { + clearTimeout(timeoutId); + this.cleanup(); + resolve(); + }); + } else { + clearTimeout(timeoutId); + resolve(); + } + }); + } +} + +// Singleton instance shared by all IPC handlers. +let queryServiceManager: QueryServiceManager | null = null; + +export function getQueryServiceManager(): QueryServiceManager { + if (!queryServiceManager) { + queryServiceManager = new QueryServiceManager(); + } + return queryServiceManager; +} + +export async function shutdownQueryService(): Promise { + if (queryServiceManager) { + await queryServiceManager.shutdown(); + queryServiceManager = null; + } +} + +/** + * Delete a dataset's search index from disk, closing it first if it is the + * one currently open in the query service. + */ +async function deleteIndex(settings: Settings, datasetId: string): Promise { + const indexDir = getIndexDir(settings, datasetId); + const manager = getQueryServiceManager(); + if (manager.openIndexDir() === indexDir) { + await manager.closeIndex(); + } + await fs.remove(indexDir); +} + +/** + * Save the current IQR model as a runnable trained pipeline: + * DIVE_Pipelines// containing .svm plus the generic-detector SVM + * template as detector.pipe. Shows up under the "trained" pipeline category. + */ +async function exportSearchModel(settings: Settings, name: string): Promise { + const safeName = name.replace(/[^a-zA-Z0-9 _-]/g, '').trim(); + if (!safeName) { + throw new Error('A model name is required'); + } + const manager = getQueryServiceManager(); + if (!manager.isReady() || !manager.openIndexDir()) { + throw new Error('No active query session to export a model from'); + } + const template = npath.join(settings.viamePath, 'configs', 'pipelines', ExportTemplateName); + if (!(await fs.pathExists(template))) { + throw new Error(`Missing pipeline template: ${template}`); + } + const outputDir = npath.join(settings.dataPath, 'DIVE_Pipelines', safeName); + await fs.ensureDir(outputDir); + await manager.exportModel(npath.join(outputDir, `${safeName}.svm`)); + await fs.copy(template, npath.join(outputDir, 'detector.pipe')); + return outputDir; +} + +export { + isVideoSearchInstalled, + getIndexDir, + getIndexStatus, + buildIndex, + deleteIndex, + exportSearchModel, +}; diff --git a/client/platform/desktop/background.ts b/client/platform/desktop/background.ts index c95e1ad59..a21a9cdd3 100644 --- a/client/platform/desktop/background.ts +++ b/client/platform/desktop/background.ts @@ -6,6 +6,7 @@ import os from 'os'; import path from 'path'; import { closeAll as closeChildren } from './backend/native/processManager'; +import { shutdownQueryService } from './backend/native/videoSearch'; import { listen, close as closeServer } from './backend/server'; import ipcListen from './backend/ipcService'; @@ -87,6 +88,9 @@ protocol.registerSchemesAsPrivileged([ async function cleanup() { closeServer(); + // Graceful first: lets the query service stop its embedded postgres + // before closeChildren() kills whatever is left. + await shutdownQueryService().catch(() => undefined); await closeChildren(); app.quit(); } diff --git a/client/platform/desktop/constants.ts b/client/platform/desktop/constants.ts index 5b24868ff..7ad06bfa8 100644 --- a/client/platform/desktop/constants.ts +++ b/client/platform/desktop/constants.ts @@ -171,6 +171,7 @@ export enum JobType { ExportTrainedPipeline, RunPipeline, RunTraining, + BuildSearchIndex, } export interface JobArgs { @@ -220,7 +221,30 @@ export interface ConversionArgs extends JobArgs { mediaList: [string, string][]; } -export type Job = ConversionArgs | RunPipeline | RunTraining | ExportTrainedPipeline; +/** Build a video search / IQR descriptor index over a dataset. */ +export interface BuildSearchIndex extends JobArgs { + type: JobType.BuildSearchIndex; + datasetId: string; + // detections: index around generic object proposals + // tracking: index around tracked proposals + // existing: index around this dataset's existing annotations + method: 'detections' | 'tracking' | 'existing'; +} + +/** Sidecar metadata written alongside a built search index. */ +export interface SearchIndexMeta { + version: number; + method: BuildSearchIndex['method']; + // frame rate the media was indexed at (video only; must match dataset fps) + fps: number; + // dataset ids covered by this index (single entry today; keyed this way so + // multi-dataset indexes remain possible later) + datasets: string[]; + createdAt: string; +} + +export type Job = ConversionArgs | RunPipeline | RunTraining + | ExportTrainedPipeline | BuildSearchIndex; export interface DesktopJob { // key unique identifier for this job @@ -232,7 +256,7 @@ export interface DesktopJob { // title whatever humans should see this job called title: string; // arguments to creation - args: RunPipeline | RunTraining | ExportTrainedPipeline | ConversionArgs; + args: RunPipeline | RunTraining | ExportTrainedPipeline | ConversionArgs | BuildSearchIndex; // datasetIds of the involved datasets datasetIds: string[]; // pid of the process spawned From 2ffba937f4ed4f49057b45607e0570e8a763ed9c Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 6 Jul 2026 21:55:58 -0400 Subject: [PATCH 2/3] Add Video Search side panel with IQR refinement loop New in-viewer context panel (desktop only) for index-backed video search and rapid model generation: build/rebuild/delete the dataset's search index (via the GPU job queue), query from the selected annotation, an external image file, or a saved .svm model warm-start, adjudicate ranked results +/- with cropped chip thumbnails and click-to-seek, refine iteratively, and save good models as runnable trained pipelines. Shared request/response types live in apispec so web can implement later. Co-Authored-By: Claude Fable 5 --- client/dive-common/apispec.ts | 54 ++ client/platform/desktop/backend/ipcService.ts | 17 +- .../desktop/backend/native/videoSearch.ts | 31 +- client/platform/desktop/constants.ts | 14 +- client/platform/desktop/frontend/api.ts | 74 ++- .../components/VideoSearchContext.vue | 467 ++++++++++++++++++ .../frontend/components/ViewerLoader.vue | 34 +- .../platform/desktop/frontend/store/jobs.ts | 4 +- .../frontend/store/queues/asyncGpuJobQueue.ts | 19 +- .../desktop/frontend/useVideoSearch.ts | 253 ++++++++++ 10 files changed, 940 insertions(+), 27 deletions(-) create mode 100644 client/platform/desktop/frontend/components/VideoSearchContext.vue create mode 100644 client/platform/desktop/frontend/useVideoSearch.ts diff --git a/client/dive-common/apispec.ts b/client/dive-common/apispec.ts index 55fca437c..7d13ae43a 100644 --- a/client/dive-common/apispec.ts +++ b/client/dive-common/apispec.ts @@ -418,6 +418,60 @@ export interface SegmentationStatusResponse { ready?: boolean; } +/** + * Video Search / IQR (rapid model generation) Types + */ + +export type VideoSearchIndexMethod = 'detections' | 'tracking' | 'existing'; + +/** Sidecar metadata describing a built search index. */ +export interface VideoSearchIndexMeta { + version: number; + method: VideoSearchIndexMethod; + // frame rate the media was indexed at (video only; must match dataset fps) + fps: number; + // dataset ids covered by this index (single entry today; keyed this way so + // multi-dataset indexes remain possible later) + datasets: string[]; + createdAt: string; +} + +export interface VideoSearchIndexStatus { + exists: boolean; + built: boolean; + meta?: VideoSearchIndexMeta; +} + +export interface VideoSearchTrackState { + frame: number; + bbox?: [number, number, number, number]; +} + +export interface VideoSearchResultTrack { + id: number; + states: VideoSearchTrackState[]; +} + +/** One ranked similarity result returned by the query service. */ +export interface VideoSearchResult { + instance_id: number; + query_id: string; + stream_id: string; + relevancy_score: number; + start_frame: number | null; + end_frame: number | null; + tracks: VideoSearchResultTrack[]; +} + +export interface VideoSearchQueryResponse { + success: boolean; + error?: string; + descriptor_count?: number; + model_available?: boolean; + results?: VideoSearchResult[]; + feedback_requests?: VideoSearchResult[]; +} + export { provideApi, useApi, diff --git a/client/platform/desktop/backend/ipcService.ts b/client/platform/desktop/backend/ipcService.ts index c50cba838..56864a666 100644 --- a/client/platform/desktop/backend/ipcService.ts +++ b/client/platform/desktop/backend/ipcService.ts @@ -507,9 +507,17 @@ export default function register() { return manager.formulateQuery(args.imagePath, args.boxes); }); - ipcMain.handle('video-search-query', async (_, args: { threshold?: number; iqrModelB64?: string }) => { + ipcMain.handle('video-search-query', async ( + _, + args: { threshold?: number; iqrModelB64?: string; iqrModelPath?: string }, + ) => { const manager = videoSearch.getQueryServiceManager(); - return manager.processQuery(args.threshold, args.iqrModelB64); + let modelB64 = args.iqrModelB64; + if (!modelB64 && args.iqrModelPath) { + // Warm-start from a saved .svm file on disk + modelB64 = (await fs.promises.readFile(args.iqrModelPath)).toString('base64'); + } + return manager.processQuery(args.threshold, modelB64); }); ipcMain.handle('video-search-refine', async (_, args: { positiveIds: number[]; negativeIds: number[] }) => { @@ -527,4 +535,9 @@ export default function register() { await manager.closeIndex(); return { success: true }; }); + + ipcMain.handle('video-search-extract-frame', async ( + _, + args: { videoPath: string; frameNum: number; fps: number }, + ) => videoSearch.extractVideoFrame(args.videoPath, args.frameNum, args.fps)); } diff --git a/client/platform/desktop/backend/native/videoSearch.ts b/client/platform/desktop/backend/native/videoSearch.ts index 94f97ab10..d48eaf54c 100644 --- a/client/platform/desktop/backend/native/videoSearch.ts +++ b/client/platform/desktop/backend/native/videoSearch.ts @@ -26,10 +26,13 @@ import { Settings, DesktopJob, DesktopJobUpdater, SearchIndexMeta, BuildSearchIndex, } from 'platform/desktop/constants'; +import type { VideoSearchIndexStatus } from 'dive-common/apispec'; import { serialize } from 'platform/desktop/backend/serializers/viame'; import { observeChild } from './processManager'; import * as common from './common'; -import { jobFileEchoMiddleware, createCustomWorkingDirectory } from './utils'; +import { + jobFileEchoMiddleware, createCustomWorkingDirectory, getBinaryPath, spawnResult, +} from './utils'; import linux from './linux'; import win32 from './windows'; @@ -67,11 +70,7 @@ function getIndexDir(settings: Settings, datasetId: string): string { return npath.join(settings.dataPath, 'DIVE_Projects', datasetId, SearchIndexFolderName); } -export interface SearchIndexStatus { - exists: boolean; - built: boolean; - meta?: SearchIndexMeta; -} +export type SearchIndexStatus = VideoSearchIndexStatus; async function getIndexStatus(settings: Settings, datasetId: string): Promise { const indexDir = getIndexDir(settings, datasetId); @@ -206,6 +205,25 @@ async function buildIndex( return jobBase; } +/** + * Extract a single video frame as a PNG for use as a query exemplar + * (the exemplar descriptor pipeline reads image files, not videos). + */ +async function extractVideoFrame(videoPath: string, frameNum: number, fps: number): Promise { + const outDir = npath.join(OS.tmpdir(), 'dive-video-search'); + await fs.ensureDir(outDir); + const outPath = npath.join(outDir, `frame_${frameNum}.png`); + const seconds = frameNum / (fps || 1); + const ffmpegPath = getBinaryPath('ffmpeg-ffprobe-static/ffmpeg'); + const result = await spawnResult(ffmpegPath, [ + '-y', '-ss', seconds.toString(), '-i', videoPath, '-frames:v', '1', outPath, + ]); + if (result.exitCode !== 0 || !(await fs.pathExists(outPath))) { + throw new Error(`Unable to extract video frame ${frameNum}: ${result.error}`); + } + return outPath; +} + /** Loose shape of a JSON response line from the Python query service. */ interface ServiceResponse { id: string; @@ -599,4 +617,5 @@ export { buildIndex, deleteIndex, exportSearchModel, + extractVideoFrame, }; diff --git a/client/platform/desktop/constants.ts b/client/platform/desktop/constants.ts index 7ad06bfa8..2489e3670 100644 --- a/client/platform/desktop/constants.ts +++ b/client/platform/desktop/constants.ts @@ -1,6 +1,7 @@ import type { DatasetMeta, DatasetMetaMutable, DatasetType, Pipe, SubType, MediaImportResponse, PipelineParams, + VideoSearchIndexMeta, VideoSearchIndexMethod, } from 'dive-common/apispec'; import { Attribute } from 'vue-media-annotator/use/AttributeTypes'; import { AttributeTrackFilter } from 'vue-media-annotator/AttributeTrackFilterControls'; @@ -228,20 +229,11 @@ export interface BuildSearchIndex extends JobArgs { // detections: index around generic object proposals // tracking: index around tracked proposals // existing: index around this dataset's existing annotations - method: 'detections' | 'tracking' | 'existing'; + method: VideoSearchIndexMethod; } /** Sidecar metadata written alongside a built search index. */ -export interface SearchIndexMeta { - version: number; - method: BuildSearchIndex['method']; - // frame rate the media was indexed at (video only; must match dataset fps) - fps: number; - // dataset ids covered by this index (single entry today; keyed this way so - // multi-dataset indexes remain possible later) - datasets: string[]; - createdAt: string; -} +export type SearchIndexMeta = VideoSearchIndexMeta; export type Job = ConversionArgs | RunPipeline | RunTraining | ExportTrainedPipeline | BuildSearchIndex; diff --git a/client/platform/desktop/frontend/api.ts b/client/platform/desktop/frontend/api.ts index 4f1fb1c71..1c5c881d5 100644 --- a/client/platform/desktop/frontend/api.ts +++ b/client/platform/desktop/frontend/api.ts @@ -7,6 +7,7 @@ import type { DatasetCalibrationResult, SegmentationPredictRequest, SegmentationPredictResponse, SegmentationStatusResponse, SegmentationStereoSegmentRequest, SegmentationStereoSegmentResponse, + VideoSearchIndexStatus, VideoSearchIndexMethod, VideoSearchQueryResponse, } from 'dive-common/apispec'; import { @@ -19,7 +20,7 @@ import { RunPipeline, RunTraining, ExportTrainedPipeline, ExportDatasetArgs, ExportConfigurationArgs, ExportMulticamEverythingArgs, DesktopMediaImportResponse, ConversionArgs, JobType, - DesktopJob, + DesktopJob, BuildSearchIndex, } from 'platform/desktop/constants'; import { gpuJobQueue, cpuJobQueue } from './store/jobs'; @@ -334,6 +335,64 @@ async function segmentationIsReady(): Promise { return window.diveDesktop.invoke('segmentation-is-ready'); } +/** + * Video Search / IQR API + */ + +async function videoSearchInstalled(): Promise { + return window.diveDesktop.invoke('video-search-installed'); +} + +async function videoSearchIndexStatus(datasetId: string): Promise { + return window.diveDesktop.invoke('video-search-index-status', datasetId); +} + +/** Queue a search index build; progress arrives via job-update events. */ +function videoSearchBuildIndex(datasetId: string, method: VideoSearchIndexMethod): void { + const args: BuildSearchIndex = { type: JobType.BuildSearchIndex, datasetId, method }; + gpuJobQueue.enqueue(args); +} + +async function videoSearchDeleteIndex(datasetId: string): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-delete-index', datasetId); +} + +async function videoSearchOpenIndex(datasetId: string): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-open-index', datasetId); +} + +async function videoSearchFormulate(imagePath: string, boxes?: number[][]): Promise { + return window.diveDesktop.invoke('video-search-formulate', { imagePath, boxes }); +} + +async function videoSearchQuery( + options: { threshold?: number; iqrModelB64?: string; iqrModelPath?: string } = {}, +): Promise { + return window.diveDesktop.invoke('video-search-query', options); +} + +/** Build a local media-server URL for an arbitrary file (e.g. thumbnails). */ +async function getMediaUrl(filePath: string): Promise { + await getClient(); + return `${_baseURL}/media?path=${encodeURIComponent(filePath)}`; +} + +async function videoSearchRefine(positiveIds: number[], negativeIds: number[]): Promise { + return window.diveDesktop.invoke('video-search-refine', { positiveIds, negativeIds }); +} + +async function videoSearchExportModel(name: string): Promise<{ success: boolean; outputDir: string }> { + return window.diveDesktop.invoke('video-search-export-model', { name }); +} + +async function videoSearchClose(): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-close'); +} + +async function videoSearchExtractFrame(videoPath: string, frameNum: number, fps: number): Promise { + return window.diveDesktop.invoke('video-search-extract-frame', { videoPath, frameNum, fps }); +} + /** * Interactive Stereo API */ @@ -656,6 +715,19 @@ export { segmentationClearImage, segmentationShutdown, segmentationIsReady, + /* Video Search / IQR */ + videoSearchInstalled, + videoSearchIndexStatus, + videoSearchBuildIndex, + videoSearchDeleteIndex, + videoSearchOpenIndex, + videoSearchFormulate, + videoSearchQuery, + videoSearchRefine, + videoSearchExportModel, + videoSearchClose, + videoSearchExtractFrame, + getMediaUrl, /* Stereo APIs */ stereoEnable, stereoDisable, diff --git a/client/platform/desktop/frontend/components/VideoSearchContext.vue b/client/platform/desktop/frontend/components/VideoSearchContext.vue new file mode 100644 index 000000000..a5927eaa0 --- /dev/null +++ b/client/platform/desktop/frontend/components/VideoSearchContext.vue @@ -0,0 +1,467 @@ + + + + + diff --git a/client/platform/desktop/frontend/components/ViewerLoader.vue b/client/platform/desktop/frontend/components/ViewerLoader.vue index da5455fc4..5ba5552cc 100644 --- a/client/platform/desktop/frontend/components/ViewerLoader.vue +++ b/client/platform/desktop/frontend/components/ViewerLoader.vue @@ -32,6 +32,10 @@ import { import Export from './Export.vue'; import JobTab from './JobTab.vue'; import DatasetSourceInfo from './DatasetSourceInfo.vue'; +import VideoSearchContext from './VideoSearchContext.vue'; +import { + createVideoSearch, provideVideoSearch, VideoSearchMediaInfo, +} from '../useVideoSearch'; import { datasets } from '../store/dataset'; import { settings } from '../store/settings'; import { runningJobs } from '../store/jobs'; @@ -48,6 +52,13 @@ function joinPath(base: string, file: string): string { return `${base.replace(/[\\/]+$/, '')}${sep}${file}`; } +// Desktop-only context panel: registered here (not in the shared context +// store) so the web build does not pick it up. +context.register({ + description: 'Video Search', + component: VideoSearchContext, +}); + const buttonOptions = { outlined: true, color: 'grey lighten-1', @@ -294,9 +305,30 @@ export default defineComponent({ originalVideoFile?: string; } | null = null; + /** + * Video Search / IQR session (index-backed similarity queries). + * Media info resolves lazily once metadata loads; multicam datasets are + * not yet supported (media stays null and the panel reports unavailable). + */ + const videoSearchMedia = ref(null); + const videoSearch = createVideoSearch(props.id, () => videoSearchMedia.value); + provideVideoSearch(videoSearch); + // Initialize segmentation when component is mounted - onMounted(() => { + onMounted(async () => { initializeSegmentation(); + try { + const meta = await loadMetadata(props.id); + if (!meta.multiCamMedia) { + videoSearchMedia.value = { + type: meta.type, + fps: meta.fps, + getImagePath: buildImagePathGetter(meta), + }; + } + } catch { + // Video search stays unavailable if metadata cannot load + } }); /** diff --git a/client/platform/desktop/frontend/store/jobs.ts b/client/platform/desktop/frontend/store/jobs.ts index 9be68c3f6..9ea4720fd 100644 --- a/client/platform/desktop/frontend/store/jobs.ts +++ b/client/platform/desktop/frontend/store/jobs.ts @@ -6,6 +6,7 @@ import { reactive, } from 'vue'; import { + BuildSearchIndex, ConversionArgs, DesktopJob, DesktopJobUpdate, @@ -159,7 +160,8 @@ function removeJobFromQueue(jobArgs: JobArgs) { break; case JobType.RunPipeline: case JobType.RunTraining: - gpuJobQueue.removeJobFromQueue(jobArgs as RunPipeline | RunTraining); + case JobType.BuildSearchIndex: + gpuJobQueue.removeJobFromQueue(jobArgs as RunPipeline | RunTraining | BuildSearchIndex); break; default: break; diff --git a/client/platform/desktop/frontend/store/queues/asyncGpuJobQueue.ts b/client/platform/desktop/frontend/store/queues/asyncGpuJobQueue.ts index 01a8483ce..3923bb4bd 100644 --- a/client/platform/desktop/frontend/store/queues/asyncGpuJobQueue.ts +++ b/client/platform/desktop/frontend/store/queues/asyncGpuJobQueue.ts @@ -2,6 +2,7 @@ import { JobType, RunPipeline, RunTraining, + BuildSearchIndex, DesktopJob, } from 'platform/desktop/constants'; import AsyncJobQueue from './asyncJobQueue'; @@ -23,30 +24,38 @@ function runTrainingsAreEquivalent(a: RunTraining, b: RunTraining) { ); } -export default class AsyncGpuJobQueue extends AsyncJobQueue { - async beginJob(spec: RunPipeline | RunTraining) { +type GpuJobSpec = RunPipeline | RunTraining | BuildSearchIndex; + +export default class AsyncGpuJobQueue extends AsyncJobQueue { + async beginJob(spec: GpuJobSpec) { let newJob: DesktopJob; if (spec.type === JobType.RunPipeline) { newJob = await this.ipcRenderer.invoke('run-pipeline', spec); } else if (spec.type === JobType.RunTraining) { newJob = await this.ipcRenderer.invoke('run-training', spec); + } else if (spec.type === JobType.BuildSearchIndex) { + newJob = await this.ipcRenderer.invoke('video-search-build-index', spec); } else { throw new Error('Unsupported job arguments provided to beginJob.'); } this.processingJobs.push(newJob); } - removeJobFromQueue(removeSpec: RunPipeline | RunTraining): void { + removeJobFromQueue(removeSpec: GpuJobSpec): void { let removeSpecIndex = -1; if (removeSpec.type === JobType.RunPipeline) { - removeSpecIndex = this.jobSpecs.findIndex((spec: RunPipeline | RunTraining) => (spec.type === JobType.RunPipeline + removeSpecIndex = this.jobSpecs.findIndex((spec: GpuJobSpec) => (spec.type === JobType.RunPipeline && spec.datasetId === removeSpec.datasetId && spec.pipeline.pipe === removeSpec.pipeline.pipe )); } else if (removeSpec.type === JobType.RunTraining) { - removeSpecIndex = this.jobSpecs.findIndex((spec: RunPipeline | RunTraining) => (spec.type === JobType.RunTraining + removeSpecIndex = this.jobSpecs.findIndex((spec: GpuJobSpec) => (spec.type === JobType.RunTraining && runTrainingsAreEquivalent(spec, removeSpec) )); + } else if (removeSpec.type === JobType.BuildSearchIndex) { + removeSpecIndex = this.jobSpecs.findIndex((spec: GpuJobSpec) => (spec.type === JobType.BuildSearchIndex + && spec.datasetId === removeSpec.datasetId + )); } if (removeSpecIndex > -1) { this.jobSpecs.splice(removeSpecIndex, 1); diff --git a/client/platform/desktop/frontend/useVideoSearch.ts b/client/platform/desktop/frontend/useVideoSearch.ts new file mode 100644 index 000000000..a14d4452d --- /dev/null +++ b/client/platform/desktop/frontend/useVideoSearch.ts @@ -0,0 +1,253 @@ +/** + * Video search / IQR session state shared between ViewerLoader (which owns + * dataset media details) and the VideoSearchContext side panel. + * + * Created and provided by ViewerLoader; injected by the panel. + */ +import { reactive, provide, inject } from 'vue'; +import type { + VideoSearchIndexStatus, VideoSearchIndexMethod, + VideoSearchResult, VideoSearchQueryResponse, +} from 'dive-common/apispec'; +import { + videoSearchInstalled, videoSearchIndexStatus, videoSearchBuildIndex, + videoSearchDeleteIndex, videoSearchOpenIndex, videoSearchFormulate, + videoSearchQuery, videoSearchRefine, videoSearchExportModel, + videoSearchClose, videoSearchExtractFrame, +} from 'platform/desktop/frontend/api'; + +export const VideoSearchInjectKey = 'DesktopVideoSearch'; + +export type Adjudication = 'positive' | 'negative'; + +export interface VideoSearchMediaInfo { + type: string; // 'image-sequence' | 'video' + fps: number; + getImagePath: (frameNum: number) => string; +} + +interface VideoSearchState { + installed: boolean | null; + status: VideoSearchIndexStatus | null; + sessionOpen: boolean; + /** Human-readable description of the operation in flight, or null. */ + busy: string | null; + error: string | null; + results: VideoSearchResult[]; + adjudications: Record; + /** Completed refinement iterations for the active query. */ + iteration: number; + modelAvailable: boolean; +} + +export function createVideoSearch( + datasetId: string, + getMediaInfo: () => VideoSearchMediaInfo | null, +) { + const state = reactive({ + installed: null, + status: null, + sessionOpen: false, + busy: null, + error: null, + results: [], + adjudications: {}, + iteration: 0, + modelAvailable: false, + }); + + async function refreshStatus() { + try { + if (state.installed === null) { + state.installed = await videoSearchInstalled(); + } + if (state.installed) { + state.status = await videoSearchIndexStatus(datasetId); + } + } catch (err) { + state.error = err instanceof Error ? err.message : String(err); + } + } + + /** Run an operation with busy/error bookkeeping. */ + async function guarded(label: string, op: () => Promise): Promise { + if (state.busy) { + return undefined; + } + state.busy = label; + state.error = null; + try { + return await op(); + } catch (err) { + state.error = err instanceof Error ? err.message : String(err); + return undefined; + } finally { + state.busy = null; + } + } + + function buildIndex(method: VideoSearchIndexMethod) { + // Runs through the GPU job queue; completion observed via the jobs store. + videoSearchBuildIndex(datasetId, method); + } + + async function deleteIndex() { + await guarded('Deleting index...', async () => { + await videoSearchDeleteIndex(datasetId); + state.sessionOpen = false; + resetQueryState(); + }); + await refreshStatus(); + } + + function resetQueryState() { + state.results = []; + state.adjudications = {}; + state.iteration = 0; + state.modelAvailable = false; + } + + async function ensureSession() { + if (!state.sessionOpen) { + await videoSearchOpenIndex(datasetId); + state.sessionOpen = true; + } + } + + function applyResponse(response: VideoSearchQueryResponse) { + if (response.results) { + state.results = response.results; + // Keep adjudications for instance ids that are still present so marks + // survive re-ranking across refinement rounds. + const keep: Record = {}; + response.results.forEach((r) => { + keep[r.instance_id] = state.adjudications[r.instance_id]; + }); + state.adjudications = keep; + } + if (response.model_available) { + state.modelAvailable = true; + } + } + + /** + * Resolve an exemplar image for a frame: the image file itself for image + * sequences, or an extracted frame for videos. + */ + async function exemplarImageForFrame(frameNum: number): Promise { + const media = getMediaInfo(); + if (!media) { + throw new Error('Dataset media has not finished loading'); + } + if (media.type === 'video') { + return videoSearchExtractFrame(media.getImagePath(frameNum), frameNum, media.fps); + } + const imagePath = media.getImagePath(frameNum); + if (!imagePath) { + throw new Error(`No image found for frame ${frameNum}`); + } + return imagePath; + } + + /** + * Start a new query from an image chip: exemplar image plus zero or more + * boxes. With boxes, the backend runs the similarity query in the same + * step; without them, an explicit query pass follows. + */ + async function queryFromImage(imagePath: string, boxes?: number[][], warmStartModelPath?: string) { + await guarded('Searching...', async () => { + await ensureSession(); + resetQueryState(); + let response = await videoSearchFormulate(imagePath, boxes); + if (!response.results || warmStartModelPath) { + response = await videoSearchQuery( + warmStartModelPath ? { iqrModelPath: warmStartModelPath } : {}, + ); + } + applyResponse(response); + }); + } + + /** Start a new query from a frame of this dataset + boxes. */ + async function queryFromFrame(frameNum: number, boxes: number[][], warmStartModelPath?: string) { + await guarded('Searching...', async () => { + await ensureSession(); + resetQueryState(); + const imagePath = await exemplarImageForFrame(frameNum); + let response = await videoSearchFormulate(imagePath, boxes); + if (!response.results || warmStartModelPath) { + response = await videoSearchQuery( + warmStartModelPath ? { iqrModelPath: warmStartModelPath } : {}, + ); + } + applyResponse(response); + }); + } + + function mark(instanceId: number, adjudication: Adjudication) { + if (state.adjudications[instanceId] === adjudication) { + state.adjudications = { ...state.adjudications, [instanceId]: undefined }; + } else { + state.adjudications = { ...state.adjudications, [instanceId]: adjudication }; + } + } + + async function refine() { + const positives: number[] = []; + const negatives: number[] = []; + Object.entries(state.adjudications).forEach(([id, adj]) => { + if (adj === 'positive') positives.push(Number(id)); + if (adj === 'negative') negatives.push(Number(id)); + }); + if (!positives.length && !negatives.length) { + state.error = 'Mark at least one result as correct or incorrect before refining'; + return; + } + await guarded('Refining...', async () => { + const response = await videoSearchRefine(positives, negatives); + applyResponse(response); + state.iteration += 1; + }); + } + + async function saveModel(name: string): Promise { + return guarded('Saving model...', async () => { + const { outputDir } = await videoSearchExportModel(name); + return outputDir; + }); + } + + async function closeSession() { + await guarded('Closing...', async () => { + await videoSearchClose(); + state.sessionOpen = false; + }); + } + + const context = { + datasetId, + state, + getMediaInfo, + refreshStatus, + buildIndex, + deleteIndex, + queryFromImage, + queryFromFrame, + exemplarImageForFrame, + mark, + refine, + saveModel, + closeSession, + }; + return context; +} + +export type VideoSearchContextType = ReturnType; + +export function provideVideoSearch(context: VideoSearchContextType) { + provide(VideoSearchInjectKey, context); +} + +export function useVideoSearch(): VideoSearchContextType | null { + return inject(VideoSearchInjectKey, null); +} From bf24a8e2314efc842e7bed1be95898b3da6d4f9c Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 7 Jul 2026 14:46:57 -0400 Subject: [PATCH 3/3] Search across all indexed datasets via one shared index Replace per-dataset index directories with a single shared search database (DIVE_SearchIndex): every table keys rows on a per-video stream identifier, so datasets are added, updated, and removed independently while one IQR session searches everything at once. - Add/Update ingests a dataset into the shared database as a job (initializing it on first use), re-using stream identifiers derived from the dataset id (image sequences) or video filename stem. - Remove deletes the dataset's rows through the query service and drops it from the membership metadata; deleting a dataset from DIVE also removes it from the index automatically. - Query results attribute back to their source dataset via stream_id; the panel labels cross-dataset results, seeks within the current dataset, crops thumbnails from each result's own media, and can filter the display to the current dataset. Co-Authored-By: Claude Fable 5 --- client/dive-common/apispec.ts | 45 ++- client/platform/desktop/backend/ipcService.ts | 39 ++- .../desktop/backend/native/videoSearch.ts | 287 +++++++++++++----- client/platform/desktop/frontend/api.ts | 26 +- .../components/VideoSearchContext.vue | 112 +++++-- .../desktop/frontend/useVideoSearch.ts | 144 +++++++-- 6 files changed, 487 insertions(+), 166 deletions(-) diff --git a/client/dive-common/apispec.ts b/client/dive-common/apispec.ts index 7d13ae43a..b63affbe2 100644 --- a/client/dive-common/apispec.ts +++ b/client/dive-common/apispec.ts @@ -424,21 +424,35 @@ export interface SegmentationStatusResponse { export type VideoSearchIndexMethod = 'detections' | 'tracking' | 'existing'; -/** Sidecar metadata describing a built search index. */ -export interface VideoSearchIndexMeta { - version: number; +/** + * One indexed media stream (video/sequence identifier) in the shared search + * database. All database rows key on this identifier, so a dataset can be + * added, updated, or removed from the index independently. + */ +export interface VideoSearchStreamEntry { + datasetId: string; method: VideoSearchIndexMethod; // frame rate the media was indexed at (video only; must match dataset fps) - fps: number; - // dataset ids covered by this index (single entry today; keyed this way so - // multi-dataset indexes remain possible later) - datasets: string[]; + fps?: number; createdAt: string; } +/** Sidecar metadata describing the shared search index. */ +export interface VideoSearchIndexMeta { + version: number; + // stream identifier (as reported in query results) -> source dataset + streams: Record; +} + export interface VideoSearchIndexStatus { - exists: boolean; + /** The shared index exists and is queryable (ITQ files present). */ built: boolean; + /** This dataset has been ingested into the index. */ + indexed: boolean; + /** The stream entry for this dataset, when indexed. */ + stream?: VideoSearchStreamEntry & { streamName: string }; + /** Total datasets in the shared index. */ + datasetCount: number; meta?: VideoSearchIndexMeta; } @@ -454,6 +468,12 @@ export interface VideoSearchResultTrack { /** One ranked similarity result returned by the query service. */ export interface VideoSearchResult { + /** Unique reference for adjudication: ":" */ + ref: string; + /** Federated session index this result came from */ + session: number; + /** Index directory this result came from */ + index_dir: string; instance_id: number; query_id: string; stream_id: string; @@ -463,6 +483,15 @@ export interface VideoSearchResult { tracks: VideoSearchResultTrack[]; } +/** One dataset present in the shared search index. */ +export interface VideoSearchIndexInfo { + /** Stream identifier query results report for this dataset */ + streamName: string; + datasetId: string; + /** Dataset display name */ + name: string; +} + export interface VideoSearchQueryResponse { success: boolean; error?: string; diff --git a/client/platform/desktop/backend/ipcService.ts b/client/platform/desktop/backend/ipcService.ts index 56864a666..f5807164e 100644 --- a/client/platform/desktop/backend/ipcService.ts +++ b/client/platform/desktop/backend/ipcService.ts @@ -199,6 +199,11 @@ export default function register() { ipcMain.handle('delete-dataset', async (event, { datasetId }: { datasetId: string }) => { const ret = await common.deleteDataset(settings.get(), datasetId); + // Also drop the dataset's rows from the shared search index (best + // effort; the index tolerates stale entries if this fails). + videoSearch.removeFromIndex(settings.get(), datasetId).catch((err) => { + console.error(`Failed to remove ${datasetId} from the search index:`, err); + }); return ret; }); @@ -478,28 +483,36 @@ export default function register() { const updater = (update: DesktopJobUpdate) => { event.sender.send('job-update', update); }; - // A rebuild invalidates whatever the service currently has open. + // The build job needs the shared database to itself (and re-ingesting + // invalidates whatever the open session has cached in memory). const manager = videoSearch.getQueryServiceManager(); - if (manager.openIndexDir() === videoSearch.getIndexDir(settings.get(), args.datasetId)) { - await manager.closeIndex(); - } + await manager.closeIndex(); return videoSearch.buildIndex(settings.get(), args, updater); }); - ipcMain.handle('video-search-delete-index', async (_, datasetId: string) => { - await videoSearch.deleteIndex(settings.get(), datasetId); + ipcMain.handle('video-search-remove-index', async (_, datasetId: string) => { + await videoSearch.removeFromIndex(settings.get(), datasetId); return { success: true }; }); - ipcMain.handle('video-search-open-index', async (_, datasetId: string) => { + ipcMain.handle('video-search-delete-index', async () => { + await videoSearch.deleteEntireIndex(settings.get()); + return { success: true }; + }); + + ipcMain.handle('video-search-list-indexes', async () => ( + videoSearch.listIndexedDatasets(settings.get()) + )); + + ipcMain.handle('video-search-open-index', async () => { const currentSettings = settings.get(); - const status = await videoSearch.getIndexStatus(currentSettings, datasetId); - if (!status.built) { - throw new Error('No built search index exists for this dataset'); + const streams = await videoSearch.listIndexedDatasets(currentSettings); + if (!streams.length) { + throw new Error('The search index is empty; add a dataset to it first'); } const manager = videoSearch.getQueryServiceManager(); - await manager.openIndex(currentSettings, videoSearch.getIndexDir(currentSettings, datasetId)); - return { success: true }; + await manager.openIndexes(currentSettings, [videoSearch.getIndexDir(currentSettings)]); + return { success: true, streams }; }); ipcMain.handle('video-search-formulate', async (_, args: { imagePath: string; boxes?: number[][] }) => { @@ -520,7 +533,7 @@ export default function register() { return manager.processQuery(args.threshold, modelB64); }); - ipcMain.handle('video-search-refine', async (_, args: { positiveIds: number[]; negativeIds: number[] }) => { + ipcMain.handle('video-search-refine', async (_, args: { positiveIds: string[]; negativeIds: string[] }) => { const manager = videoSearch.getQueryServiceManager(); return manager.refine(args.positiveIds, args.negativeIds); }); diff --git a/client/platform/desktop/backend/native/videoSearch.ts b/client/platform/desktop/backend/native/videoSearch.ts index d48eaf54c..ea9b37cfe 100644 --- a/client/platform/desktop/backend/native/videoSearch.ts +++ b/client/platform/desktop/backend/native/videoSearch.ts @@ -1,17 +1,19 @@ /** * Video Search / IQR backend for Desktop * - * Provides per-dataset search index management (build via a process_video.py - * job, status, delete) and a persistent query service manager - * (viame.core.query_service) that holds the KWIVER query/IQR pipeline open so - * refinement iterations are interactive. One index may be open at a time (the - * embedded PostgreSQL instance binds a fixed port), enforced here by closing - * the previous index before opening another. + * Manages ONE shared search index covering every indexed dataset, plus a + * persistent query service manager (viame.core.query_service) that holds the + * KWIVER query/IQR pipeline open so refinement iterations are interactive. * - * On-disk layout (per dataset): - * /DIVE_Projects//search_index/ - * index_meta.json - written by DIVE on successful build - * ingest_list.txt - media list handed to process_video.py + * All database tables key rows on a per-video stream identifier + * (VIDEO_NAME), so datasets are added to, updated in, and removed from the + * shared database independently; a query searches every indexed dataset in + * one IQR session. + * + * On-disk layout: + * /DIVE_SearchIndex/ + * index_meta.json - stream -> dataset membership, written by DIVE + * .txt - media list handed to process_video.py * database/ - embedded PostgreSQL data + ITQ/LSH index files */ @@ -24,9 +26,9 @@ import { EventEmitter } from 'events'; import { Settings, DesktopJob, DesktopJobUpdater, - SearchIndexMeta, BuildSearchIndex, + SearchIndexMeta, BuildSearchIndex, JsonMeta, } from 'platform/desktop/constants'; -import type { VideoSearchIndexStatus } from 'dive-common/apispec'; +import type { VideoSearchIndexStatus, VideoSearchIndexInfo } from 'dive-common/apispec'; import { serialize } from 'platform/desktop/backend/serializers/viame'; import { observeChild } from './processManager'; import * as common from './common'; @@ -36,7 +38,7 @@ import { import linux from './linux'; import win32 from './windows'; -const SearchIndexFolderName = 'search_index'; +const GlobalIndexFolderName = 'DIVE_SearchIndex'; const IndexMetaFileName = 'index_meta.json'; const QueryPipelineName = 'query_retrieval_and_iqr.pipe'; const ExportTemplateName = npath.join('templates', 'detector_generic_svm.pipe'); @@ -66,29 +68,89 @@ async function isVideoSearchInstalled(settings: Settings): Promise { return checks.every((c) => c); } -function getIndexDir(settings: Settings, datasetId: string): string { - return npath.join(settings.dataPath, 'DIVE_Projects', datasetId, SearchIndexFolderName); +function getIndexDir(settings: Settings): string { + return npath.join(settings.dataPath, GlobalIndexFolderName); +} + +/** Filesystem/SQL-safe stream identifier token. */ +function sanitizeName(name: string): string { + return name.replace(/[^a-zA-Z0-9_-]/g, '_'); } export type SearchIndexStatus = VideoSearchIndexStatus; -async function getIndexStatus(settings: Settings, datasetId: string): Promise { - const indexDir = getIndexDir(settings, datasetId); - const metaPath = npath.join(indexDir, IndexMetaFileName); +async function readIndexMeta(settings: Settings): Promise { + const metaPath = npath.join(getIndexDir(settings), IndexMetaFileName); if (!(await fs.pathExists(metaPath))) { - return { exists: await fs.pathExists(indexDir), built: false }; + return { version: 1, streams: {} }; + } + return (await fs.readJson(metaPath)) as SearchIndexMeta; +} + +async function writeIndexMeta(settings: Settings, meta: SearchIndexMeta): Promise { + await fs.writeJson(npath.join(getIndexDir(settings), IndexMetaFileName), meta, { spaces: 2 }); +} + +/** The stream identifier a dataset's rows key on in the shared database. */ +function streamNameForDataset(datasetId: string, meta: JsonMeta): string { + if (meta.type === 'video') { + // Videos are attributed by their filename stem (process_video derives + // the ingest stream name from the input file's basename) + return npath.parse(meta.originalVideoFile).name; } - const meta = (await fs.readJson(metaPath)) as SearchIndexMeta; - // Sanity: the ITQ model files must exist for the index to be queryable - const itqDir = npath.join(indexDir, 'database', 'ITQ'); - const built = (await fs.pathExists(itqDir)) + return sanitizeName(datasetId); +} + +/** Whether the shared index is queryable (ITQ/LSH files present). */ +async function indexIsBuilt(settings: Settings): Promise { + const itqDir = npath.join(getIndexDir(settings), 'database', 'ITQ'); + return (await fs.pathExists(itqDir)) && (await fs.readdir(itqDir)).some((f) => f.startsWith('itq.model')); - return { exists: true, built, meta }; +} + +async function getIndexStatus(settings: Settings, datasetId: string): Promise { + const meta = await readIndexMeta(settings); + const built = await indexIsBuilt(settings); + const streamName = Object.keys(meta.streams) + .find((s) => meta.streams[s].datasetId === datasetId); + return { + built, + indexed: built && streamName !== undefined, + stream: streamName !== undefined + ? { ...meta.streams[streamName], streamName } : undefined, + datasetCount: built + ? new Set(Object.values(meta.streams).map((s) => s.datasetId)).size : 0, + meta, + }; +} + +/** Every dataset present in the shared search index. */ +async function listIndexedDatasets(settings: Settings): Promise { + const meta = await readIndexMeta(settings); + const entries = await Promise.all( + Object.entries(meta.streams).map(async ([streamName, stream]) => { + let name = stream.datasetId; + try { + const projectInfo = await common.getValidatedProjectDir(settings, stream.datasetId); + const dsMeta = await common.loadJsonMetadata(projectInfo.metaFileAbsPath); + name = dsMeta.name || stream.datasetId; + } catch { + // Dataset may have been deleted; keep the id as the display name. + } + return { streamName, datasetId: stream.datasetId, name }; + }), + ); + return entries; } /** - * Build (or rebuild) the search index for a dataset by running - * process_video.py --init --build-index as a DIVE job. + * Add (or update) one dataset in the shared search index as a DIVE job. + * + * The job chain: ensure the shared postgres is up (or initialize it on the + * first ever build), delete the dataset's previous rows when updating, run + * the ingest pipeline via process_video.py, and refresh the ITQ/LSH index + * over the full database. The caller must close any open query session + * first (the job needs the default postgres port). */ async function buildIndex( settings: Settings, @@ -108,15 +170,23 @@ async function buildIndex( throw new Error('Search indexes are not yet supported on multi-camera datasets'); } - const indexDir = getIndexDir(settings, datasetId); - // Rebuild from scratch: the DB re-init inside process_video handles the - // database folder, but a stale meta file must not mark a failed build as - // usable. - await fs.remove(npath.join(indexDir, IndexMetaFileName)); + const indexDir = getIndexDir(settings); await fs.ensureDir(indexDir); + const indexMeta = await readIndexMeta(settings); + const streamName = streamNameForDataset(datasetId, meta); + const existing = indexMeta.streams[streamName]; + if (existing && existing.datasetId !== datasetId) { + throw new Error( + `Stream identifier '${streamName}' already belongs to dataset ` + + `'${existing.datasetId}'; rename the media file to index both.`, + ); + } + // Media list: image sequences list every frame; videos list the one file. - const ingestList = npath.join(indexDir, 'ingest_list.txt'); + // The list file is named after the stream so image-sequence results + // attribute back to the dataset directly. + const ingestList = npath.join(indexDir, `${sanitizeName(streamName)}.txt`); if (meta.type === 'video') { const videoAbsPath = npath.join(meta.originalBasePath, meta.originalVideoFile); await fs.writeFile(ingestList, `${videoAbsPath}\n`); @@ -129,47 +199,74 @@ async function buildIndex( const viameConstants = platform.getViameConstants(settings); const pythonExe = platform.getViamePythonExe(settings); - const processVideo = npath.join(settings.viamePath, 'configs', 'process_video.py'); + const configsDir = npath.join(settings.viamePath, 'configs'); + const firstBuild = !(await fs.pathExists(npath.join(indexDir, 'database', 'SQL'))); + + const command: string[] = [viameConstants.setupScriptAbs]; + if (!firstBuild) { + // Adding to an existing database: make sure it is running (the caller + // closed the query service, so the default port is free)... + command.push(`"${pythonExe}" "${npath.join(configsDir, 'database_tool.py')}" start`); + if (existing) { + // ...and drop the dataset's previous rows before re-ingest + const removeSql = npath.join(indexDir, `${sanitizeName(streamName)}_remove.sql`); + await fs.writeFile(removeSql, [ + 'DELETE FROM TRACK_DESCRIPTOR_TRACK WHERE UID IN ' + + `(SELECT UID FROM TRACK_DESCRIPTOR WHERE VIDEO_NAME = '${streamName}');`, + 'DELETE FROM TRACK_DESCRIPTOR_HISTORY WHERE UID IN ' + + `(SELECT UID FROM TRACK_DESCRIPTOR WHERE VIDEO_NAME = '${streamName}');`, + `DELETE FROM TRACK_DESCRIPTOR WHERE VIDEO_NAME = '${streamName}';`, + `DELETE FROM DESCRIPTOR WHERE VIDEO_NAME = '${streamName}';`, + `DELETE FROM OBJECT_TRACK WHERE VIDEO_NAME = '${streamName}';`, + '', + ].join('\n')); + command.push( + `psql -h localhost -d postgres -v ON_ERROR_STOP=1 -f "${removeSql}"`, + ); + } + } - const command = [ - `${viameConstants.setupScriptAbs} &&`, - `"${pythonExe}" "${processVideo}"`, - '--init --no-reset-prompt', + const processVideoInvocation = [ + `"${pythonExe}" "${npath.join(configsDir, 'process_video.py')}"`, + firstBuild ? '--init' : '', + '--no-reset-prompt', `-l "${ingestList}"`, `-p "pipelines/${IndexPipelines[method]}"`, '-o database', '--build-index', `-install "${settings.viamePath}"`, - ]; + ].filter((part) => part.length); if (meta.type === 'video') { - command.push(`-frate ${meta.fps}`); + processVideoInvocation.push(`-frate ${meta.fps}`); } // Index around this dataset's existing annotations if (method === 'existing') { - const detectionsCsv = npath.join(indexDir, 'input_detections.csv'); + const detectionsCsv = npath.join(indexDir, `${sanitizeName(streamName)}_detections.csv`); const csvStream = fs.createWriteStream(detectionsCsv); const inputData = await common.loadAnnotationFile(projectInfo.trackFileAbsPath); await serialize(csvStream, inputData, meta); csvStream.end(); - command.push(`-id "${detectionsCsv}"`); + processVideoInvocation.push(`-id "${detectionsCsv}"`); } + command.push(processVideoInvocation.join(' ')); + const fullCommand = command.join(' && '); const jobWorkDir = await createCustomWorkingDirectory(settings, 'search_index', datasetId.replace('/', '_')); const joblog = npath.join(jobWorkDir, 'runlog.txt'); - const job = observeChild(spawn(command.join(' '), { + const job = observeChild(spawn(fullCommand, { shell: viameConstants.shell, cwd: indexDir, })); const jobBase: DesktopJob = { key: `search_index_${job.pid}_${jobWorkDir}`, - command: command.join(' '), + command: fullCommand, jobType: 'pipeline', pid: job.pid, args, - title: `Build search index (${method})`, + title: `${existing ? 'Update' : 'Add'} search index (${method})`, workingDir: jobWorkDir, datasetIds: [datasetId], exitCode: job.exitCode, @@ -185,15 +282,15 @@ async function buildIndex( job.on('exit', async (code) => { if (code === 0) { - const indexMeta: SearchIndexMeta = { - version: 1, - method, - fps: meta.fps, - datasets: [datasetId], - createdAt: (new Date()).toISOString(), - }; try { - await fs.writeJson(npath.join(indexDir, IndexMetaFileName), indexMeta, { spaces: 2 }); + const latest = await readIndexMeta(settings); + latest.streams[streamName] = { + datasetId, + method, + fps: meta.type === 'video' ? meta.fps : undefined, + createdAt: (new Date()).toISOString(), + }; + await writeIndexMeta(settings, latest); } catch (err) { console.error('Failed to write search index metadata', err); } @@ -256,8 +353,8 @@ export class QueryServiceManager extends EventEmitter { private requestCounter = 0; - /** Index directory currently opened in the service, if any. */ - private currentIndexDir: string | null = null; + /** Index directories currently opened in the service (primary first). */ + private currentIndexDirs: string[] = []; // First open loads the descriptor index + pipeline; queries train SVMs. private readonly requestTimeoutMs = 300000; @@ -266,8 +363,8 @@ export class QueryServiceManager extends EventEmitter { return this.process !== null && this.process.exitCode === null; } - openIndexDir(): string | null { - return this.isReady() ? this.currentIndexDir : null; + openIndexDirs(): string[] { + return this.isReady() ? this.currentIndexDirs : []; } private generateRequestId(): string { @@ -440,18 +537,23 @@ export class QueryServiceManager extends EventEmitter { // -------------------------------------------------------------- commands - /** Open an index, closing any previously open one (one postgres at a time). */ - async openIndex(settings: Settings, indexDir: string): Promise { + /** + * Open one or more indexes for federated search (primary first), closing + * any previously open set if it differs. Each index gets its own embedded + * postgres instance on an incremented port inside the service. + */ + async openIndexes(settings: Settings, indexDirs: string[]): Promise { await this.ensureStarted(settings); - if (this.currentIndexDir === indexDir) { + if (indexDirs.length === this.currentIndexDirs.length + && indexDirs.every((dir, i) => this.currentIndexDirs[i] === dir)) { return; } - if (this.currentIndexDir) { + if (this.currentIndexDirs.length) { await this.closeIndex(); } - const response = await this.sendRequest({ command: 'open_index', index_dir: indexDir }, 'Open index'); - QueryServiceManager.check(response, 'Opening the search index'); - this.currentIndexDir = indexDir; + const response = await this.sendRequest({ command: 'open_index', index_dirs: indexDirs }, 'Open index'); + QueryServiceManager.check(response, 'Opening the search indexes'); + this.currentIndexDirs = indexDirs; } async formulateQuery(imagePath: string, boxes?: number[][]): Promise { @@ -472,7 +574,7 @@ export class QueryServiceManager extends EventEmitter { return QueryServiceManager.check(response, 'Query'); } - async refine(positiveIds: number[], negativeIds: number[]): Promise { + async refine(positiveIds: string[], negativeIds: string[]): Promise { const response = await this.sendRequest({ command: 'refine', positive_ids: positiveIds, @@ -489,9 +591,24 @@ export class QueryServiceManager extends EventEmitter { return QueryServiceManager.check(response, 'Model export'); } + /** + * Delete all database rows for the given stream identifiers in an index. + * The service adopts or temporarily starts the index's postgres; if the + * index was open, the session set is closed for consistency. + */ + async removeStreams(indexDir: string, streams: string[]): Promise { + const response = await this.sendRequest({ + command: 'remove_streams', index_dir: indexDir, streams, + }, 'Stream removal'); + QueryServiceManager.check(response, 'Stream removal'); + if (response.closed_session) { + this.currentIndexDirs = []; + } + } + async closeIndex(): Promise { if (!this.isReady()) { - this.currentIndexDir = null; + this.currentIndexDirs = []; return; } try { @@ -499,7 +616,7 @@ export class QueryServiceManager extends EventEmitter { } catch { // best effort - shutting the process down also stops postgres } - this.currentIndexDir = null; + this.currentIndexDirs = []; } // ------------------------------------------------------------- lifecycle @@ -517,7 +634,7 @@ export class QueryServiceManager extends EventEmitter { } this.process = null; - this.currentIndexDir = null; + this.currentIndexDirs = []; this.emit('shutdown'); } @@ -573,16 +690,30 @@ export async function shutdownQueryService(): Promise { } /** - * Delete a dataset's search index from disk, closing it first if it is the - * one currently open in the query service. + * Remove one dataset from the shared search index: delete its database rows + * (via the query service, which owns postgres lifecycle) and drop it from + * the membership metadata. Stale ITQ hash entries are tolerated by the + * query engine and cleaned up on the next index build. */ -async function deleteIndex(settings: Settings, datasetId: string): Promise { - const indexDir = getIndexDir(settings, datasetId); - const manager = getQueryServiceManager(); - if (manager.openIndexDir() === indexDir) { - await manager.closeIndex(); +async function removeFromIndex(settings: Settings, datasetId: string): Promise { + const meta = await readIndexMeta(settings); + const streams = Object.keys(meta.streams) + .filter((s) => meta.streams[s].datasetId === datasetId); + if (!streams.length) { + return; } - await fs.remove(indexDir); + const manager = getQueryServiceManager(); + await manager.ensureStarted(settings); + await manager.removeStreams(getIndexDir(settings), streams); + streams.forEach((s) => { delete meta.streams[s]; }); + await writeIndexMeta(settings, meta); +} + +/** Delete the entire shared search index from disk. */ +async function deleteEntireIndex(settings: Settings): Promise { + const manager = getQueryServiceManager(); + await manager.closeIndex(); + await fs.remove(getIndexDir(settings)); } /** @@ -596,7 +727,7 @@ async function exportSearchModel(settings: Settings, name: string): Promise { - return window.diveDesktop.invoke('video-search-delete-index', datasetId); +/** Delete the entire shared search index from disk. */ +async function videoSearchDeleteIndex(): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-delete-index'); } -async function videoSearchOpenIndex(datasetId: string): Promise<{ success: boolean }> { - return window.diveDesktop.invoke('video-search-open-index', datasetId); +async function videoSearchListIndexes(): Promise { + return window.diveDesktop.invoke('video-search-list-indexes'); +} + +async function videoSearchOpenIndex(): Promise<{ + success: boolean; streams: VideoSearchIndexInfo[]; +}> { + return window.diveDesktop.invoke('video-search-open-index'); +} + +/** Remove one dataset's rows from the shared search index. */ +async function videoSearchRemoveIndex(datasetId: string): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-remove-index', datasetId); } async function videoSearchFormulate(imagePath: string, boxes?: number[][]): Promise { @@ -377,7 +389,7 @@ async function getMediaUrl(filePath: string): Promise { return `${_baseURL}/media?path=${encodeURIComponent(filePath)}`; } -async function videoSearchRefine(positiveIds: number[], negativeIds: number[]): Promise { +async function videoSearchRefine(positiveIds: string[], negativeIds: string[]): Promise { return window.diveDesktop.invoke('video-search-refine', { positiveIds, negativeIds }); } @@ -719,7 +731,9 @@ export { videoSearchInstalled, videoSearchIndexStatus, videoSearchBuildIndex, + videoSearchRemoveIndex, videoSearchDeleteIndex, + videoSearchListIndexes, videoSearchOpenIndex, videoSearchFormulate, videoSearchQuery, diff --git a/client/platform/desktop/frontend/components/VideoSearchContext.vue b/client/platform/desktop/frontend/components/VideoSearchContext.vue index a5927eaa0..74e39a739 100644 --- a/client/platform/desktop/frontend/components/VideoSearchContext.vue +++ b/client/platform/desktop/frontend/components/VideoSearchContext.vue @@ -31,10 +31,18 @@ export default defineComponent({ const saveModelName = ref(''); const saveModelDialog = ref(false); /** Small cache of cropped result thumbnails keyed by instance id. */ - const thumbnails = ref>({}); + const thumbnails = ref>({}); const state = computed(() => search?.state ?? null); + /** Display filter: limit visible results to the current dataset. */ + const onlyThisDataset = ref(false); + const displayedResults = computed(() => { + const all = state.value?.results ?? []; + if (!onlyThisDataset.value || !search) return all; + return all.filter((r) => search.resultDatasetId(r) === search.datasetId); + }); + const indexBuilding = computed(() => runningJobs.value.some((item) => ( item.job.exitCode === null && search !== null @@ -114,18 +122,18 @@ export default defineComponent({ await search.queryFromImage(ret.filePaths[0], undefined, modelPath); } - async function deleteIndexConfirm() { + async function removeFromIndexConfirm() { if (!search) return; const confirmed = await prompt({ - title: 'Delete Search Index', - text: ['Delete this dataset\'s search index?', - 'The index can be rebuilt later, but building takes time.'], + title: 'Remove From Search Index', + text: ['Remove this dataset from the search index?', + 'It can be re-added later, but indexing takes time.'], confirm: true, - positiveButton: 'Delete', + positiveButton: 'Remove', negativeButton: 'Cancel', }); if (confirmed) { - await search.deleteIndex(); + await search.removeFromIndex(); } } @@ -144,12 +152,20 @@ export default defineComponent({ } function seekToResult(result: VideoSearchResult) { + // Only results from the currently open dataset can be seeked to. + if (search && search.resultDatasetId(result) !== search.datasetId) { + return; + } const target = result.tracks[0]?.states[0]?.frame ?? result.start_frame; if (target !== null && target !== undefined) { handler.seekFrame(target); } } + function isLocalResult(result: VideoSearchResult): boolean { + return search !== null && search.resultDatasetId(result) === search.datasetId; + } + function resultFrame(result: VideoSearchResult): number | null { return result.tracks[0]?.states[0]?.frame ?? result.start_frame ?? null; } @@ -160,12 +176,14 @@ export default defineComponent({ /** Crop a result chip out of its source frame into a data URL. */ async function loadThumbnail(result: VideoSearchResult) { - if (!search || thumbnails.value[result.instance_id]) return; + if (!search || thumbnails.value[result.ref]) return; const frameNum = resultFrame(result); const bbox = resultBox(result); if (frameNum === null) return; try { - const imagePath = await search.exemplarImageForFrame(frameNum); + const resultDataset = search.resultDatasetId(result); + if (!resultDataset) return; + const imagePath = await search.imageForDatasetFrame(resultDataset, frameNum); const url = await getMediaUrl(imagePath); const image = new Image(); image.src = url; @@ -185,7 +203,7 @@ export default defineComponent({ ctx.drawImage(image, x1, y1, w, h, 0, 0, canvas.width, canvas.height); thumbnails.value = { ...thumbnails.value, - [result.instance_id]: canvas.toDataURL('image/jpeg', 0.8), + [result.ref]: canvas.toDataURL('image/jpeg', 0.8), }; } } catch { @@ -210,10 +228,13 @@ export default defineComponent({ thumbnails, queryFromSelectedTrack, queryFromImageFile, - deleteIndexConfirm, + removeFromIndexConfirm, saveModel, seekToResult, resultFrame, + isLocalResult, + onlyThisDataset, + displayedResults, }; }, }); @@ -239,11 +260,17 @@ export default defineComponent({ - {{ indexBuilding ? 'building...' : (state.status && state.status.built ? 'built' : 'not built') }} + {{ indexBuilding ? 'indexing...' : (state.status && state.status.indexed ? 'indexed' : 'not indexed') }} +
+ {{ state.status.datasetCount }} dataset(s) in the shared search index +
- {{ state.status && state.status.built ? 'Rebuild' : 'Build' }} + {{ state.status && state.status.indexed ? 'Update' : 'Add to index' }} - Delete + Remove
@@ -289,7 +316,7 @@ export default defineComponent({ block class="mb-1" color="primary" - :disabled="!(state.status && state.status.built) || selectedTrackBox === null || !!state.busy" + :disabled="!(state.status && state.status.datasetCount) || selectedTrackBox === null || !!state.busy" @click="queryFromSelectedTrack" > Search from selected annotation @@ -298,7 +325,7 @@ export default defineComponent({ x-small block class="mb-1" - :disabled="!(state.status && state.status.built) || !!state.busy" + :disabled="!(state.status && state.status.datasetCount) || !!state.busy" @click="queryFromImageFile(false)" > Search from image file... @@ -306,13 +333,22 @@ export default defineComponent({ Search with saved model (.svm)... + + -