Add optional video search capabilities to DIVE [DRAFT]

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<string, VideoSearchStreamEntry>;
+}
+
+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: "<session>:<instance_id>" */
+  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,

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));
 }