diff --git a/client/dive-common/apispec.ts b/client/dive-common/apispec.ts index 55fca437c..b63affbe2 100644 --- a/client/dive-common/apispec.ts +++ b/client/dive-common/apispec.ts @@ -418,6 +418,89 @@ export interface SegmentationStatusResponse { ready?: boolean; } +/** + * Video Search / IQR (rapid model generation) Types + */ + +export type VideoSearchIndexMethod = 'detections' | 'tracking' | 'existing'; + +/** + * 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; + 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 { + /** 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; +} + +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 { + /** 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; + relevancy_score: number; + start_frame: number | null; + end_frame: number | null; + 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; + 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 91a5664d1..f5807164e 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, @@ -197,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; }); @@ -459,4 +466,91 @@ 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); + }; + // 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(); + await manager.closeIndex(); + return videoSearch.buildIndex(settings.get(), args, updater); + }); + + ipcMain.handle('video-search-remove-index', async (_, datasetId: string) => { + await videoSearch.removeFromIndex(settings.get(), datasetId); + return { success: true }; + }); + + 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 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.openIndexes(currentSettings, [videoSearch.getIndexDir(currentSettings)]); + return { success: true, streams }; + }); + + 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; iqrModelPath?: string }, + ) => { + const manager = videoSearch.getQueryServiceManager(); + 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: string[]; negativeIds: string[] }) => { + 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 }; + }); + + 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 new file mode 100644 index 000000000..ea9b37cfe --- /dev/null +++ b/client/platform/desktop/backend/native/videoSearch.ts @@ -0,0 +1,754 @@ +/** + * Video Search / IQR backend for Desktop + * + * 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. + * + * 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 + */ + +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, JsonMeta, +} from 'platform/desktop/constants'; +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'; +import { + jobFileEchoMiddleware, createCustomWorkingDirectory, getBinaryPath, spawnResult, +} from './utils'; +import linux from './linux'; +import win32 from './windows'; + +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'); + +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): 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 readIndexMeta(settings: Settings): Promise { + const metaPath = npath.join(getIndexDir(settings), IndexMetaFileName); + if (!(await fs.pathExists(metaPath))) { + 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; + } + 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')); +} + +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; +} + +/** + * 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, + 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); + 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. + // 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`); + } 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 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 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') { + processVideoInvocation.push(`-frate ${meta.fps}`); + } + + // Index around this dataset's existing annotations + if (method === 'existing') { + 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(); + 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(fullCommand, { + shell: viameConstants.shell, + cwd: indexDir, + })); + + const jobBase: DesktopJob = { + key: `search_index_${job.pid}_${jobWorkDir}`, + command: fullCommand, + jobType: 'pipeline', + pid: job.pid, + args, + title: `${existing ? 'Update' : 'Add'} 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) { + try { + 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); + } + } + updater({ + ...jobBase, body: [''], exitCode: code, endTime: new Date(), + }); + }); + 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; + 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 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; + + isReady(): boolean { + return this.process !== null && this.process.exitCode === null; + } + + openIndexDirs(): string[] { + return this.isReady() ? this.currentIndexDirs : []; + } + + 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 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 (indexDirs.length === this.currentIndexDirs.length + && indexDirs.every((dir, i) => this.currentIndexDirs[i] === dir)) { + return; + } + if (this.currentIndexDirs.length) { + await this.closeIndex(); + } + 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 { + 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: string[], negativeIds: string[]): 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'); + } + + /** + * 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.currentIndexDirs = []; + return; + } + try { + await this.sendRequest({ command: 'close_index' }, 'Close index'); + } catch { + // best effort - shutting the process down also stops postgres + } + this.currentIndexDirs = []; + } + + // ------------------------------------------------------------- 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.currentIndexDirs = []; + 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; + } +} + +/** + * 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 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; + } + 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)); +} + +/** + * 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.openIndexDirs().length) { + 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, + listIndexedDatasets, + buildIndex, + removeFromIndex, + deleteEntireIndex, + exportSearchModel, + extractVideoFrame, +}; 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..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'; @@ -171,6 +172,7 @@ export enum JobType { ExportTrainedPipeline, RunPipeline, RunTraining, + BuildSearchIndex, } export interface JobArgs { @@ -220,7 +222,21 @@ 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: VideoSearchIndexMethod; +} + +/** Sidecar metadata written alongside a built search index. */ +export type SearchIndexMeta = VideoSearchIndexMeta; + +export type Job = ConversionArgs | RunPipeline | RunTraining + | ExportTrainedPipeline | BuildSearchIndex; export interface DesktopJob { // key unique identifier for this job @@ -232,7 +248,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 diff --git a/client/platform/desktop/frontend/api.ts b/client/platform/desktop/frontend/api.ts index 4f1fb1c71..1b3d03dc9 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, VideoSearchIndexInfo, } 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,76 @@ 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); +} + +/** Delete the entire shared search index from disk. */ +async function videoSearchDeleteIndex(): Promise<{ success: boolean }> { + return window.diveDesktop.invoke('video-search-delete-index'); +} + +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 { + 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: string[], negativeIds: string[]): 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 +727,21 @@ export { segmentationClearImage, segmentationShutdown, segmentationIsReady, + /* Video Search / IQR */ + videoSearchInstalled, + videoSearchIndexStatus, + videoSearchBuildIndex, + videoSearchRemoveIndex, + videoSearchDeleteIndex, + videoSearchListIndexes, + 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..74e39a739 --- /dev/null +++ b/client/platform/desktop/frontend/components/VideoSearchContext.vue @@ -0,0 +1,519 @@ + + + + + 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..9aaf8f1b5 --- /dev/null +++ b/client/platform/desktop/frontend/useVideoSearch.ts @@ -0,0 +1,333 @@ +/** + * 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. + * + * All indexed datasets share one search database, so every query searches + * every indexed dataset in a single IQR session. Each result's stream_id + * maps back to its source dataset via the index membership metadata. + */ +import { reactive, provide, inject } from 'vue'; +import type { + VideoSearchIndexStatus, VideoSearchIndexMethod, VideoSearchIndexInfo, + VideoSearchResult, VideoSearchQueryResponse, +} from 'dive-common/apispec'; +import { + videoSearchInstalled, videoSearchIndexStatus, videoSearchBuildIndex, + videoSearchRemoveIndex, videoSearchOpenIndex, + videoSearchFormulate, videoSearchQuery, videoSearchRefine, + videoSearchExportModel, videoSearchClose, videoSearchExtractFrame, + loadMetadata, +} 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; + /** stream identifier -> source dataset info for the open index. */ + streams: Record; + /** Human-readable description of the operation in flight, or null. */ + busy: string | null; + error: string | null; + results: VideoSearchResult[]; + /** Keyed by result ref (":"). */ + adjudications: Record; + /** Completed refinement iterations for the active query. */ + iteration: number; + modelAvailable: boolean; +} + +/** Renderer-safe path join (node 'path' is unavailable here). */ +function isAbsolutePath(p: string): boolean { + return /^([A-Za-z]:[\\/]|[\\/])/.test(p); +} +function joinPath(base: string, file: string): string { + if (!base) return file; + const sep = base.includes('\\') ? '\\' : '/'; + return `${base.replace(/[\\/]+$/, '')}${sep}${file}`; +} + +export function createVideoSearch( + datasetId: string, + getMediaInfo: () => VideoSearchMediaInfo | null, +) { + const state = reactive({ + installed: null, + status: null, + sessionOpen: false, + streams: {}, + busy: null, + error: null, + results: [], + adjudications: {}, + iteration: 0, + modelAvailable: false, + }); + + /** Media info per dataset for cross-dataset result display. */ + const mediaInfoCache: Record = {}; + + async function getMediaInfoFor(dsId: string): Promise { + if (dsId === datasetId) { + return getMediaInfo(); + } + if (!(dsId in mediaInfoCache)) { + try { + const meta = await loadMetadata(dsId); + const { + originalBasePath, originalImageFiles, type, originalVideoFile, fps, + } = meta; + mediaInfoCache[dsId] = { + type, + fps, + getImagePath: (frameNum: number): string => { + if (type === 'video') { + return joinPath(originalBasePath, originalVideoFile || ''); + } + const imagePath = originalImageFiles?.[frameNum]; + if (!imagePath) return ''; + return isAbsolutePath(imagePath) + ? imagePath : joinPath(originalBasePath, imagePath); + }, + }; + } catch { + mediaInfoCache[dsId] = null; + } + } + return mediaInfoCache[dsId]; + } + + 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); + } + + /** Remove this dataset from the shared search index. */ + async function removeFromIndex() { + await guarded('Removing from index...', async () => { + await videoSearchRemoveIndex(datasetId); + state.sessionOpen = false; + state.streams = {}; + resetQueryState(); + }); + await refreshStatus(); + } + + function resetQueryState() { + state.results = []; + state.adjudications = {}; + state.iteration = 0; + state.modelAvailable = false; + } + + /** + * Open the shared search index (covering every indexed dataset). The + * backend is a no-op when it is already open. + */ + async function ensureSession() { + const { streams } = await videoSearchOpenIndex(); + const map: Record = {}; + streams.forEach((info) => { map[info.streamName] = info; }); + state.streams = map; + state.sessionOpen = true; + } + + /** The dataset a result belongs to (via its stream identifier). */ + function resultDatasetId(result: VideoSearchResult): string | null { + return state.streams[result.stream_id]?.datasetId ?? null; + } + + /** Display name of a result's dataset, or null when it is this dataset. */ + function resultDatasetName(result: VideoSearchResult): string | null { + const info = state.streams[result.stream_id]; + if (!info) return null; + return info.datasetId === datasetId ? null : (info.name || info.datasetId); + } + + function applyResponse(response: VideoSearchQueryResponse) { + if (response.results) { + state.results = response.results; + // Keep adjudications for refs that are still present so marks + // survive re-ranking across refinement rounds. + const keep: Record = {}; + response.results.forEach((r) => { + keep[r.ref] = state.adjudications[r.ref]; + }); + state.adjudications = keep; + } + if (response.model_available) { + state.modelAvailable = true; + } + } + + /** + * Resolve an image on disk for a frame of any open dataset: the image + * file itself for image sequences, or an extracted frame for videos. + */ + async function imageForDatasetFrame(dsId: string, frameNum: number): Promise { + const media = await getMediaInfoFor(dsId); + if (!media) { + throw new Error(`Media for dataset ${dsId} is unavailable`); + } + 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; + } + + async function exemplarImageForFrame(frameNum: number): Promise { + return imageForDatasetFrame(datasetId, frameNum); + } + + /** + * Start a new query from an image chip: exemplar image plus zero or more + * boxes. Descriptors are computed once (on this dataset's primary + * session) and the similarity query fans out to every open index. + */ + 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(ref: string, adjudication: Adjudication) { + if (state.adjudications[ref] === adjudication) { + state.adjudications = { ...state.adjudications, [ref]: undefined }; + } else { + state.adjudications = { ...state.adjudications, [ref]: adjudication }; + } + } + + async function refine() { + const positives: string[] = []; + const negatives: string[] = []; + Object.entries(state.adjudications).forEach(([ref, adj]) => { + if (adj === 'positive') positives.push(ref); + if (adj === 'negative') negatives.push(ref); + }); + 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; + state.streams = {}; + }); + } + + const context = { + datasetId, + state, + getMediaInfo, + getMediaInfoFor, + refreshStatus, + buildIndex, + removeFromIndex, + queryFromImage, + queryFromFrame, + exemplarImageForFrame, + imageForDatasetFrame, + resultDatasetId, + resultDatasetName, + 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); +}