diff --git a/client/platform/desktop/backend/native/common.spec.ts b/client/platform/desktop/backend/native/common.spec.ts index fe0a9eff3..4f470a8e5 100644 --- a/client/platform/desktop/backend/native/common.spec.ts +++ b/client/platform/desktop/backend/native/common.spec.ts @@ -7,7 +7,7 @@ import { Console } from 'console'; import type { ConversionArgs, DesktopJob, - DesktopJobUpdate, DesktopJobUpdater, JsonMeta, Settings, + DesktopJobUpdate, DesktopJobUpdater, JsonMeta, RunTraining, Settings, } from 'platform/desktop/constants'; import { Attribute } from 'viame-web-common/apispec'; @@ -136,6 +136,25 @@ mockfs({ '/home/user/viamedata': { // eslint-disable-next-line @typescript-eslint/camelcase DIVE_Jobs: { + goodTrainingJob: { + // eslint-disable-next-line @typescript-eslint/camelcase + category_models: { + 'detector.pipe': '', + 'trained_detector.zip': '', + }, + }, + badTrainingJob: { + missingModelFolder: {}, + }, + missingPipeTrainingJob: { + // eslint-disable-next-line @typescript-eslint/camelcase + category_models: { + 'trained_detector.zip': '', + }, + }, + }, + // eslint-disable-next-line @typescript-eslint/camelcase + DIVE_Pipelines: { /* Empty */ }, // eslint-disable-next-line @typescript-eslint/camelcase @@ -212,8 +231,6 @@ mockfs({ }), 'result_whatever.json': JSON.stringify({}), auxiliary: {}, - - }, }, }, @@ -228,7 +245,7 @@ describe('native.common', () => { expect(pipes.detector.pipes).toHaveLength(4); expect(pipes.tracker.pipes).toHaveLength(5); expect(pipes.generate.pipes).toHaveLength(4); - expect(pipes.training).toBeUndefined(); + expect(pipes.trained).toBeUndefined(); }); it('getValidatedProjectDir loads correct project directory', async () => { @@ -322,6 +339,49 @@ describe('native.common', () => { expect(meta.type).toBe('video'); }); + it('processing good Trained Pipeline folder', async () => { + const trainingArgs: RunTraining = { + datasetIds: ['randomID'], + pipelineName: 'trainedPipelineName', + trainingConfig: 'trainingConfig', + }; + const contents = await common.processTrainedPipeline(settings, trainingArgs, '/home/user/viamedata/DIVE_Jobs/goodTrainingJob/'); + expect(contents).toEqual(['detector.pipe', 'trained_detector.zip']); + //Data should be moved out of the current folder + const sourceFolder = fs.readdirSync('/home/user/viamedata/DIVE_Jobs/goodTrainingJob/category_models'); + expect(sourceFolder.length).toBe(0); + //Folders hould be created for new pipeline + const pipelineFolder = '/home/user/viamedata/DIVE_Pipelines/trainedPipelineName'; + const exists = fs.existsSync(pipelineFolder); + expect(exists).toBe(true); + const folderContents = fs.readdirSync(pipelineFolder); + expect(folderContents.length).toBe(2); + }); + + it('processing bad Trained Pipeline folders', async () => { + const trainingArgs: RunTraining = { + datasetIds: ['randomID'], + pipelineName: 'trainedBadPipelineName', + trainingConfig: 'trainingConfig', + }; + expect(common.processTrainedPipeline(settings, trainingArgs, '/home/user/viamedata/DIVE_Jobs/badTrainingJob/')).rejects.toThrow( + 'Path: /home/user/viamedata/Dive_Jobs/badTrainingJob/category_models does not exist', + ); + expect(common.processTrainedPipeline(settings, trainingArgs, '/home/user/viamedata/DIVE_Jobs/missingPipeTrainingJob/')).rejects.toThrow( + 'Could not located trained pipe file inside of /home/user/viamedata/Dive_Jobs/missingPipeTrainingJob/category_models', + ); + }); + + it('getPipelineList lists pipelines with Trained pipelines', async () => { + const exists = await fs.pathExists(settings.viamePath); + expect(exists).toBe(true); + const pipes = await common.getPipelineList(settings); + expect(pipes).toBeTruthy(); + expect(pipes.detector.pipes).toHaveLength(4); + expect(pipes.tracker.pipes).toHaveLength(5); + expect(pipes.generate.pipes).toHaveLength(4); + expect(pipes.trained.pipes).toHaveLength(1); + }); it('getAtributes', async () => { const meta = await common.getAttributes(settings, 'metaAttributesID'); diff --git a/client/platform/desktop/backend/native/common.ts b/client/platform/desktop/backend/native/common.ts index 61cf6f2fb..5720e6190 100644 --- a/client/platform/desktop/backend/native/common.ts +++ b/client/platform/desktop/backend/native/common.ts @@ -16,14 +16,16 @@ import * as viameSerializers from 'platform/desktop/backend/serializers/viame'; import { websafeImageTypes, websafeVideoTypes, otherImageTypes, otherVideoTypes, - JsonMeta, Settings, JsonMetaCurrentVersion, DesktopMetadata, - DesktopJobUpdater, ConvertMedia, Attributes, + JsonMeta, Settings, JsonMetaCurrentVersion, DesktopMetadata, DesktopJobUpdater, + ConvertMedia, RunTraining, Attributes, } from 'platform/desktop/constants'; import processTrackAttributes from './attributeProcessor'; import { cleanString, makeid } from './utils'; const ProjectsFolderName = 'DIVE_Projects'; const JobsFolderName = 'DIVE_Jobs'; +const PipelinesFolderName = 'DIVE_Pipelines'; + const AuxFolderName = 'auxiliary'; const JsonTrackFileName = /^result(_.*)?\.json$/; @@ -214,6 +216,41 @@ async function getPipelineList(settings: Settings): Promise { }; } }); + + // Now lets add to it the trained pipelines by recursively looking in the dir + const allowedTrainedPatterns = /^detector.+|^tracker.+|^generate.+|^trained_detector\.zip|^trained_tracker\.zip|^trained_generate\.zip/; + const trainedPipelinePath = npath.join(settings.dataPath, PipelinesFolderName); + const trainedExists = await fs.pathExists(trainedPipelinePath); + if (!trainedExists) return ret; + const trainedPipeFolders = await fs.readdir(trainedPipelinePath); + await Promise.all(trainedPipeFolders.map(async (item) => { + const pipeFolder = npath.join(trainedPipelinePath, item); + const pipeFolderExists = await fs.pathExists(pipeFolder); + if (!pipeFolderExists) return false; + let pipesInFolder = await fs.readdir(pipeFolder); + pipesInFolder = pipesInFolder.filter( + (p: string) => p.match(allowedTrainedPatterns) && !p.match(disallowedPatterns), + ); + if (pipesInFolder.length >= 2) { + const pipeName = pipesInFolder.find((pipe) => pipe && pipe.indexOf('.pipe') !== -1); + if (pipeName) { + const pipeInfo = { + name: item, + type: 'trained', + pipe: npath.join(pipeFolder, pipeName), + }; + if ('trained' in ret) { + ret.trained.pipes.push(pipeInfo); + } else { + ret.trained = { + pipes: [pipeInfo], + description: 'trained pipes', + }; + } + } + } + return true; + })); return ret; } @@ -413,6 +450,40 @@ async function processOtherAnnotationFiles( } return { fps, processedFiles, attributes }; } +/** + * 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, '/category_models'); + const exists = await fs.pathExists(trainedDir); + if (!exists) { + throw new Error(`Path: ${trainedDir} does not exist`); + } + const folderContents = await fs.readdir(trainedDir); + const pipes = folderContents.filter((p) => p.match(allowedPatterns)); + + 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); + const destpath = npath.join(folderName, item); + await fs.move(abspath, destpath, { overwrite: true }); + })); + return folderContents; +} async function _initializeAppDataDir(settings: Settings) { await fs.ensureDir(settings.dataPath); @@ -645,6 +716,7 @@ export { saveDetections, saveMetadata, completeConversion, + processTrainedPipeline, getAttributes, setAttribute, deleteAttribute, diff --git a/client/platform/desktop/backend/native/viame.ts b/client/platform/desktop/backend/native/viame.ts index 84a1b9bf2..a1e8279c1 100644 --- a/client/platform/desktop/backend/native/viame.ts +++ b/client/platform/desktop/backend/native/viame.ts @@ -46,7 +46,10 @@ async function runPipeline( throw new Error(isValid); } - const pipelinePath = npath.join(settings.viamePath, PipelineRelativeDir, pipeline.pipe); + let pipelinePath = npath.join(settings.viamePath, PipelineRelativeDir, pipeline.pipe); + if (runPipelineArgs.pipeline.type === 'trained') { + pipelinePath = pipeline.pipe; + } const projectInfo = await common.getValidatedProjectDir(settings, datasetId); const meta = await common.loadJsonMetadata(projectInfo.metaFileAbsPath); const jobWorkDir = await common.createKwiverRunWorkingDir(settings, [meta], pipeline.name); @@ -218,6 +221,7 @@ async function train( `--config "${configFilePath}"`, '--no-query', '--no-adv-prints', + '--no-embedded-pipe', ]; const job = spawn(command.join(' '), { @@ -249,11 +253,27 @@ async function train( job.stdout.on('data', jobFileEchoMiddleware(jobBase, updater, joblog)); job.stderr.on('data', jobFileEchoMiddleware(jobBase, updater, joblog)); - job.on('exit', (code) => { + job.on('exit', async (code) => { + let exitCode = code; + const bodyText = ['']; + if (code === 0) { + try { + await common.processTrainedPipeline( + settings, runTrainingArgs, jobWorkDir, + ); + } catch (err) { + console.error(err); + exitCode = 1; + bodyText.unshift(err.toString('utf-8')); + fs.appendFile(joblog, bodyText[0], (error) => { + if (error) throw error; + }); + } + } updater({ ...jobBase, - body: [''], - exitCode: code, + body: bodyText, + exitCode, endTime: new Date(), }); }); diff --git a/client/platform/desktop/frontend/components/MultiTrainingMenu.vue b/client/platform/desktop/frontend/components/MultiTrainingMenu.vue index aa7b83499..ad682b02b 100644 --- a/client/platform/desktop/frontend/components/MultiTrainingMenu.vue +++ b/client/platform/desktop/frontend/components/MultiTrainingMenu.vue @@ -2,15 +2,35 @@ import type { DataTableHeader } from 'vuetify'; import { - computed, defineComponent, onBeforeMount, set, del, reactive, + computed, defineComponent, onBeforeMount, set, del, reactive, ref, } from '@vue/composition-api'; -import { DatasetMeta, TrainingConfigs, useApi } from 'viame-web-common/apispec'; +import { + DatasetMeta, Pipelines, TrainingConfigs, useApi, +} from 'viame-web-common/apispec'; + import { datasets } from '../store/dataset'; export default defineComponent({ setup(_, { root }) { const { getTrainingConfigurations, runTraining } = useApi(); + const { getPipelineList } = useApi(); + const unsortedPipelines = ref({} as Pipelines); + onBeforeMount(async () => { + unsortedPipelines.value = await getPipelineList(); + }); + + const trainedPipelines = computed(() => { + if (unsortedPipelines.value.trained) { + return unsortedPipelines.value.trained.pipes.map((item) => item.name); + } + return []; + }); + + const nameRules = [ + (val: string) => (!trainedPipelines.value.includes(val) || 'A Trained pipeline with that name already exists'), + ]; + const data = reactive({ stagedItems: {} as Record, trainingOutputName: '', @@ -95,6 +115,7 @@ export default defineComponent({ toggleStaged, isReadyToTrain, runTrainingOnFolder, + nameRules, available: { items: availableItems, headers: headersTmpl.concat( @@ -138,9 +159,9 @@ export default defineComponent({