diff --git a/client/dive-common/apispec.ts b/client/dive-common/apispec.ts index 977006bdd..7372daa16 100644 --- a/client/dive-common/apispec.ts +++ b/client/dive-common/apispec.ts @@ -895,6 +895,24 @@ export interface VideoSearchQueryResponse { feedback_requests?: VideoSearchResult[]; } +/** One result placed in descriptor space around the query (at the origin). */ +export interface VideoSearchLayoutPoint { + ref: string; + /** Offset from the query along the three principal axes of the result set. */ + position: [number, number, number]; + /** Full-dimensional descriptor distance from the query. */ + distance: number; +} + +export interface VideoSearchLayoutResponse { + success: boolean; + error?: string; + dimensions?: number; + points?: VideoSearchLayoutPoint[]; + /** Results with no stored descriptor vector, left out of the layout. */ + missing?: string[]; +} + export { provideApi, useApi, diff --git a/client/dive-common/review/chipRenderer.spec.ts b/client/dive-common/review/chipRenderer.spec.ts index b7a77cada..c500eabd9 100644 --- a/client/dive-common/review/chipRenderer.spec.ts +++ b/client/dive-common/review/chipRenderer.spec.ts @@ -1,5 +1,5 @@ import { - chipRegion, chipScale, chipSizeFor, frameRegion, toChipPoint, toImagePoint, + boxRegion, chipRegion, chipScale, chipSizeFor, frameRegion, toChipPoint, toImagePoint, } from './chipRenderer'; describe('chipRegion', () => { @@ -30,6 +30,18 @@ describe('chipRegion', () => { }); }); +describe('boxRegion', () => { + it('is the box itself, with no context and no squaring', () => { + expect(boxRegion([10, 20, 30, 60])).toEqual({ + x: 10, y: 20, width: 20, height: 40, + }); + expect(boxRegion([30, 60, 10, 20])).toEqual({ + x: 10, y: 20, width: 20, height: 40, + }); + expect(boxRegion([5, 5, 5, 5])).toMatchObject({ width: 1, height: 1 }); + }); +}); + describe('chipSizeFor', () => { it('rounds cell sizes up to a bucket and caps at the largest', () => { expect(chipSizeFor(100)).toBe(128); diff --git a/client/dive-common/review/chipRenderer.ts b/client/dive-common/review/chipRenderer.ts index a842e82bc..a6f33ba8a 100644 --- a/client/dive-common/review/chipRenderer.ts +++ b/client/dive-common/review/chipRenderer.ts @@ -16,6 +16,11 @@ export interface ChipRenderOptions { aspect?: number; /** Box outline colour; omit to draw no outline. */ outline?: string; + /** + * Crop to the box itself: no context and no squaring, so the chip takes the + * box's own shape. Overrides padding and aspect. + */ + tight?: boolean; /** JPEG quality. */ quality?: number; } @@ -71,6 +76,17 @@ export function chipRegion(bounds: RectBounds, padding: number, aspect = 1): Chi }; } +/** Exactly the box, whatever its shape. */ +export function boxRegion(bounds: RectBounds): ChipRegion { + const [x1, y1, x2, y2] = bounds; + return { + x: Math.min(x1, x2), + y: Math.min(y1, y2), + width: Math.max(1, Math.abs(x2 - x1)), + height: Math.max(1, Math.abs(y2 - y1)), + }; +} + /** Pixel sizes chips are rendered at; cells pick the smallest that covers them. */ export const CHIP_SIZE_BUCKETS = [128, 192, 256, 384, 512, 768]; @@ -108,9 +124,11 @@ export function renderChip( bounds: RectBounds | null, options: ChipRenderOptions, ): RenderedChip { - const region = bounds - ? chipRegion(bounds, options.padding, options.aspect) - : frameRegion(frame.width, frame.height, options.aspect); + let region = frameRegion(frame.width, frame.height, options.aspect); + if (bounds) { + region = options.tight + ? boxRegion(bounds) : chipRegion(bounds, options.padding, options.aspect); + } const scale = chipScale(region, options.size); const canvas = document.createElement('canvas'); canvas.width = Math.max(1, Math.round(region.width * scale)); diff --git a/client/dive-common/review/chipStore.ts b/client/dive-common/review/chipStore.ts index be19fbba4..db61d7f55 100644 --- a/client/dive-common/review/chipStore.ts +++ b/client/dive-common/review/chipStore.ts @@ -19,6 +19,8 @@ export interface ChipStoreOptions { /** Chip width / height, matching the cells the chips are shown in. */ aspect: number; outline: string; + /** Crop to the box itself, ignoring padding and aspect. */ + tight?: boolean; } export interface ChipStoreDeps { @@ -72,7 +74,8 @@ export function createChipStore(deps: ChipStoreDeps, initial: ChipStoreOptions) /** Re-render everything when the crop or resolution changes. */ function setOptions(next: ChipStoreOptions) { if (next.padding === options.padding && next.size === options.size - && next.aspect === options.aspect && next.outline === options.outline) { + && next.aspect === options.aspect && next.outline === options.outline + && !!next.tight === !!options.tight) { return; } options = { ...next }; diff --git a/client/platform/desktop/backend/ipcService.ts b/client/platform/desktop/backend/ipcService.ts index bbde40635..7dcb3b259 100644 --- a/client/platform/desktop/backend/ipcService.ts +++ b/client/platform/desktop/backend/ipcService.ts @@ -717,6 +717,11 @@ export default function register() { return manager.refine(args.positiveIds, args.negativeIds); }); + ipcMain.handle('video-search-layout', async (_, args: { refs: string[] }) => { + const manager = videoSearch.getQueryServiceManager(); + return manager.layoutResults(args.refs); + }); + ipcMain.handle('video-search-export-model', async (_, args: { name: string }) => { const outputDir = await videoSearch.exportSearchModel(settings.get(), args.name); return { success: true, outputDir }; diff --git a/client/platform/desktop/backend/native/common.spec.ts b/client/platform/desktop/backend/native/common.spec.ts index bc78fd90a..9149386c5 100644 --- a/client/platform/desktop/backend/native/common.spec.ts +++ b/client/platform/desktop/backend/native/common.spec.ts @@ -3196,12 +3196,12 @@ describe('resumable training jobs', () => { expect(found.map((j) => j.workingDir)).toEqual([npath.join(jobsDir, 'interrupted')]); }); - it('excludes legacy successful runs with an emptied category_models', async () => { + it('excludes legacy successful runs with an emptied trained_model', async () => { mockJobsFolder({ legacySuccess: jobDirConfig( 'legacySuccess', {}, - ['deep_training', 'category_models', 'input_folder_list.txt', 'input_truth_list.txt'], + ['deep_training', 'trained_model', 'input_folder_list.txt', 'input_truth_list.txt'], ), legacyInterrupted: jobDirConfig('legacyInterrupted', {}), }); diff --git a/client/platform/desktop/backend/native/common.ts b/client/platform/desktop/backend/native/common.ts index 02d13c896..b86e6c4a1 100644 --- a/client/platform/desktop/backend/native/common.ts +++ b/client/platform/desktop/backend/native/common.ts @@ -77,6 +77,7 @@ import processTrackAttributes from './attributeProcessor'; import { upgrade } from './migrations'; // TODO: Check to Refactor this import { getMultiCamUrls, transcodeMultiCam } from './multiCamUtils'; +import { extractModelPackTo } from './modelPack'; import { loadRegistrationFiles, referenceCameraName, saveRegistrationToDatasetDir, } from './cameraRegistration'; @@ -1904,13 +1905,45 @@ async function ingestDataFiles( processedFiles, annotationsSaved, meta, warnings: outwarnings, speciesStyling, }; } +const TrainedModelPackName = 'trained_model.zip'; +// Newest first; the legacy name is what releases before the pack wrote +const TrainedModelFolderNames = ['trained_model', 'category_models']; + +async function findTrainedModelDir(workingDir: string): Promise { + const candidates = TrainedModelFolderNames.map((name) => npath.join(workingDir, name)); + const present = await Promise.all(candidates.map((candidate) => fs.pathExists(candidate))); + return candidates.find((_, index) => present[index]); +} + /** * Need to take the trained pipeline if it exists and place it in the DIVE_Pipelines folder */ async function processTrainedPipeline(settings: Settings, args: RunTraining, workingDir: string) { - //Look for trained_detector.zip and detector.pipe and move them to DIVE_Pipelines folder const allowedPatterns = /^detector.+|^tracker.+|^generate.+/; - const trainedDir = npath.join(workingDir, '/trained_model'); + const baseFolder = npath.join(settings.dataPath, PipelinesFolderName); + if (!fs.existsSync(baseFolder)) { + await fs.mkdir(baseFolder); + } + const folderName = npath.join(baseFolder, args.pipelineName); + if (!fs.existsSync(folderName)) { + await fs.mkdir(folderName); + } + + // Current VIAME packs the whole output into one zip + const pack = npath.join(workingDir, TrainedModelPackName); + if (await fs.pathExists(pack)) { + await extractModelPackTo(pack, folderName); + const packContents = await fs.readdir(folderName); + if (!packContents.some((p) => p.match(allowedPatterns))) { + throw new Error(`Could not located trained pipe file inside of ${pack}`); + } + await fs.remove(pack); + return packContents; + } + + // An unpacked folder: SVM trainers keep one, as did older VIAME releases + const trainedDir = await findTrainedModelDir(workingDir) + ?? npath.join(workingDir, TrainedModelFolderNames[0]); const exists = await fs.pathExists(trainedDir); if (!exists) { throw new Error(`Path: ${trainedDir} does not exist`); @@ -1921,15 +1954,6 @@ async function processTrainedPipeline(settings: Settings, args: RunTraining, wor if (!pipes.length) { throw new Error(`Could not located trained pipe file inside of ${trainedDir}`); } - const baseFolder = npath.join(settings.dataPath, PipelinesFolderName); - if (!fs.existsSync(baseFolder)) { - await fs.mkdir(baseFolder); - } - - const folderName = npath.join(baseFolder, args.pipelineName); - if (!fs.existsSync(folderName)) { - await fs.mkdir(folderName); - } //Move detector and model to the new folder await Promise.all(folderContents.map(async (item) => { const abspath = npath.join(trainedDir, item); @@ -1977,10 +2001,10 @@ async function findResumableTrainingJobs(settings: Settings): Promise fs.existsSync(npath.join(workingDir, f)))) return; // Manifests predating final-status recording never carry an end time; a - // successful run's models were moved out of category_models, leaving it empty + // successful run's models were moved out of the model folder, leaving it empty if (manifest.endTime === undefined) { - const modelsDir = npath.join(workingDir, 'category_models'); - if (fs.existsSync(modelsDir) && (await fs.readdir(modelsDir)).length === 0) return; + const modelsDir = await findTrainedModelDir(workingDir); + if (modelsDir && (await fs.readdir(modelsDir)).length === 0) return; } // The jobs folder may have been relocated since the manifest was written results.push({ ...manifest, workingDir }); diff --git a/client/platform/desktop/backend/native/modelPack.spec.ts b/client/platform/desktop/backend/native/modelPack.spec.ts index 258e52d28..6ecf1b9ea 100644 --- a/client/platform/desktop/backend/native/modelPack.spec.ts +++ b/client/platform/desktop/backend/native/modelPack.spec.ts @@ -4,7 +4,10 @@ import fs from 'fs-extra'; import archiver from 'archiver'; import type { Pipe } from 'dive-common/apispec'; import { PipelinesFolderName, Settings } from 'platform/desktop/constants'; -import { exportModelPack, importModelPack, modelPackPaths } from './modelPack'; +import * as common from './common'; +import { + exportModelPack, extractModelPackTo, importModelPack, modelPackPaths, +} from './modelPack'; const layouts: { name: string; files: string[]; expected: string[] }[] = JSON.parse( fs.readFileSync(path.resolve(__dirname, '../../../../../testutils/model-pack-layouts.json'), 'utf8'), @@ -31,6 +34,27 @@ async function makeZip(files: string[], destination: string) { }); } +it('extracts a training pack as it is', async () => { + const source = path.join(temp, 'trained_model.zip'); + await makeZip(['detector.pipe', 'trained_detector.pth', 'test_results/plots/pr.csv'], source); + const written = await extractModelPackTo(source, path.join(temp, 'out')); + expect(written.sort()).toEqual(['detector.pipe', 'test_results/plots/pr.csv', 'trained_detector.pth']); + expect(await fs.readFile(path.join(temp, 'out', 'test_results', 'plots', 'pr.csv'), 'utf8')).toBe('contents of test_results/plots/pr.csv'); +}); + +it('moves a trained pack into the pipelines folder and drops the archive', async () => { + const jobDir = path.join(temp, 'DIVE_Jobs', 'job'); + await fs.ensureDir(jobDir); + await makeZip(['detector.pipe', 'trained_detector.pth', 'MODEL_CARD.md'], path.join(jobDir, 'trained_model.zip')); + const contents = await common.processTrainedPipeline(settings, { pipelineName: 'fish' } as never, jobDir); + expect(contents.sort()).toEqual(['MODEL_CARD.md', 'detector.pipe', 'trained_detector.pth']); + expect(await fs.pathExists(path.join(jobDir, 'trained_model.zip'))).toBe(false); + expect(await fs.pathExists(path.join(temp, PipelinesFolderName, 'fish', 'detector.pipe'))).toBe(true); + await makeZip(['notes.txt'], path.join(jobDir, 'trained_model.zip')); + await expect(common.processTrainedPipeline(settings, { pipelineName: 'nopipe' } as never, jobDir)) + .rejects.toThrow('Could not located trained pipe'); +}); + it.each(layouts)('imports $name layout with paths and contents intact', async ({ files, expected }) => { expect([...modelPackPaths(files).values()]).toEqual(expected); const source = path.join(temp, 'fish.zip'); diff --git a/client/platform/desktop/backend/native/modelPack.ts b/client/platform/desktop/backend/native/modelPack.ts index 4958cb34b..014b2c166 100644 --- a/client/platform/desktop/backend/native/modelPack.ts +++ b/client/platform/desktop/backend/native/modelPack.ts @@ -121,6 +121,38 @@ export async function importModelPack(settings: Settings, filename: string): Pro } } +/** + * Unpack an archive as it is, every entry under destination. Used for the + * pack `viame train` writes, whose layout is already the pipeline folder's. + */ +export async function extractModelPackTo(filename: string, destination: string): Promise { + let zip: yauzl.ZipFile | undefined; + const written: string[] = []; + try { + zip = await openArchive(filename); + const entries = (await readEntries(zip)).filter((entry) => !entry.fileName.endsWith('/')); + for (const entry of entries) { + const relative = entry.fileName.replace(/\\/g, '/'); + const parts = relative.split('/'); + if (relative.startsWith('/') || parts.some((p) => p === '..' || p === '' || p.includes(':'))) { + throw new Error(`Unsafe model archive path: ${entry.fileName}`); + } + const target = path.join(destination, ...parts); + await fs.ensureDir(path.dirname(target)); + const source = await new Promise((resolve, reject) => { + zip!.openReadStream(entry, (error, stream) => { + if (error || !stream) reject(error || new Error('Unable to read ZIP entry')); else resolve(stream); + }); + }); + await streamPipeline(source, fs.createWriteStream(target)); + written.push(relative); + } + } finally { + zip?.close(); + } + return written; +} + export async function exportModelPack(settings: Settings, model: Pipe, destination: string): Promise { const root = await fs.realpath(path.join(settings.dataPath, PipelinesFolderName)); const folder = await fs.realpath(path.dirname(model.pipe)); diff --git a/client/platform/desktop/backend/native/videoSearch.ts b/client/platform/desktop/backend/native/videoSearch.ts index 8301bf8fd..1e892ad1c 100644 --- a/client/platform/desktop/backend/native/videoSearch.ts +++ b/client/platform/desktop/backend/native/videoSearch.ts @@ -770,6 +770,14 @@ export class QueryServiceManager extends EventEmitter { return QueryServiceManager.check(response, 'Query refinement'); } + async layoutResults(refs: string[]): Promise { + const response = await this.sendRequest({ + command: 'layout_results', + refs, + }, 'Result layout'); + return QueryServiceManager.check(response, 'Result layout'); + } + async exportModel(outputPath?: string): Promise { const response = await this.sendRequest({ command: 'export_model', diff --git a/client/platform/desktop/frontend/api.ts b/client/platform/desktop/frontend/api.ts index 74d87a1aa..f0e57d99b 100644 --- a/client/platform/desktop/frontend/api.ts +++ b/client/platform/desktop/frontend/api.ts @@ -12,6 +12,7 @@ import type { PipelineJobResult, ScoringDatasetSummary, ScoringJobArgs, ScoringResult, ScoringResultSummary, ScoringSourceOptions, VideoSearchIndexStatus, VideoSearchIndexMethod, VideoSearchQueryResponse, VideoSearchIndexInfo, + VideoSearchLayoutResponse, } from 'dive-common/apispec'; import axios, { AxiosInstance } from 'axios'; import { watch } from 'vue'; @@ -686,6 +687,11 @@ async function videoSearchRefine(positiveIds: string[], negativeIds: string[]): return window.diveDesktop.invoke('video-search-refine', { positiveIds, negativeIds }); } +/** Place the given results of the open query in 3D descriptor space around it. */ +async function videoSearchLayout(refs: string[]): Promise { + return window.diveDesktop.invoke('video-search-layout', { refs }); +} + async function videoSearchExportModel(name: string): Promise<{ success: boolean; outputDir: string }> { return window.diveDesktop.invoke('video-search-export-model', { name }); } @@ -1148,6 +1154,7 @@ export { videoSearchFormulate, videoSearchQuery, videoSearchRefine, + videoSearchLayout, videoSearchExportModel, videoSearchClose, videoSearchExtractFrame, diff --git a/client/platform/desktop/frontend/components/QueryExemplar.vue b/client/platform/desktop/frontend/components/QueryExemplar.vue index b392c9a1a..03337bfb9 100644 --- a/client/platform/desktop/frontend/components/QueryExemplar.vue +++ b/client/platform/desktop/frontend/components/QueryExemplar.vue @@ -100,20 +100,23 @@ export default defineComponent({