initial push
This commit is contained in:
Philipp
176
backend/services/generate-json-yolox.js
Normal file
176
backend/services/generate-json-yolox.js
Normal file
@@ -0,0 +1,176 @@
|
||||
const TrainingProject = require('../models/TrainingProject.js');
|
||||
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js')
|
||||
const LabelStudioProject = require('../models/LabelStudioProject.js')
|
||||
const Annotation = require('../models/Annotation.js')
|
||||
const Images = require('../models/Images.js')
|
||||
const fs = require('fs');
|
||||
|
||||
|
||||
async function generateTrainingJson(trainingId){
|
||||
// trainingId is now project_details_id
|
||||
const trainingProjectDetails = await TrainingProjectDetails.findByPk(trainingId);
|
||||
if (!trainingProjectDetails) throw new Error('No TrainingProjectDetails found for project_details_id ' + trainingId);
|
||||
const detailsObj = trainingProjectDetails.get({ plain: true });
|
||||
// Get parent project for name
|
||||
const trainingProject = await TrainingProject.findByPk(detailsObj.project_id);
|
||||
// Get split percentages (assume they are stored as train_percent, valid_percent, test_percent)
|
||||
const trainPercent = detailsObj.train_percent || 85;
|
||||
const validPercent = detailsObj.valid_percent || 10;
|
||||
const testPercent = detailsObj.test_percent || 5;
|
||||
|
||||
let cocoImages = [];
|
||||
let cocoAnnotations = [];
|
||||
let cocoCategories = [];
|
||||
let categoryMap = {};
|
||||
let categoryId = 0;
|
||||
let imageid = 0;
|
||||
let annotationid = 0;
|
||||
|
||||
for (const cls of detailsObj.class_map) {
|
||||
const asgMap = [];
|
||||
const listAsg = cls[1];
|
||||
for(const asg of listAsg){
|
||||
asgMap.push ({ original: asg[0], mapped: asg[1] });
|
||||
// Build category list and mapping
|
||||
if (asg[1] && !(asg[1] in categoryMap)) {
|
||||
categoryMap[asg[1]] = categoryId;
|
||||
cocoCategories.push({ id: categoryId, name: asg[1], supercategory: '' });
|
||||
categoryId++;
|
||||
}
|
||||
}
|
||||
const images = await Images.findAll({ where: { project_id: cls[0] } });
|
||||
for(const image of images){
|
||||
imageid += 1;
|
||||
let fileName = image.image_path;
|
||||
if (fileName.includes('%20')) {
|
||||
fileName = fileName.replace(/%20/g, ' ');
|
||||
}
|
||||
if (fileName && fileName.startsWith('/data/local-files/?d=')) {
|
||||
fileName = fileName.replace('/data/local-files/?d=', '');
|
||||
fileName = fileName.replace('/home/kitraining/home/kitraining/', '');
|
||||
}
|
||||
if (fileName && fileName.startsWith('home/kitraining/To_Annotate/')) {
|
||||
fileName = fileName.replace('home/kitraining/To_Annotate/','');
|
||||
}
|
||||
// Get annotations for this image
|
||||
const annotations = await Annotation.findAll({ where: { image_id: image.image_id } });
|
||||
// Use image.width and image.height from DB (populated from original_width/original_height)
|
||||
cocoImages.push({
|
||||
id: imageid,
|
||||
file_name: fileName,
|
||||
width: image.width || 0,
|
||||
height: image.height || 0
|
||||
});
|
||||
for (const annotation of annotations) {
|
||||
// Translate class name using asgMap
|
||||
let mappedClass = annotation.Label;
|
||||
for (const mapEntry of asgMap) {
|
||||
if (annotation.Label === mapEntry.original) {
|
||||
mappedClass = mapEntry.mapped;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Only add annotation if mappedClass is valid
|
||||
if (mappedClass && mappedClass in categoryMap) {
|
||||
annotationid += 1;
|
||||
let area = 0;
|
||||
if (annotation.width && annotation.height) {
|
||||
area = annotation.width * annotation.height;
|
||||
}
|
||||
cocoAnnotations.push({
|
||||
id: annotationid,
|
||||
image_id: imageid,
|
||||
category_id: categoryMap[mappedClass],
|
||||
bbox: [annotation.x, annotation.y, annotation.width, annotation.height],
|
||||
area: area,
|
||||
iscrowd: annotation.iscrowd || 0
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Shuffle images for random split using seed
|
||||
function seededRandom(seed) {
|
||||
let x = Math.sin(seed++) * 10000;
|
||||
return x - Math.floor(x);
|
||||
}
|
||||
function shuffle(array, seed) {
|
||||
for (let i = array.length - 1; i > 0; i--) {
|
||||
const j = Math.floor(seededRandom(seed + i) * (i + 1));
|
||||
[array[i], array[j]] = [array[j], array[i]];
|
||||
}
|
||||
}
|
||||
// Use seed from detailsObj if present, else default to 42
|
||||
const splitSeed = detailsObj.seed !== undefined && detailsObj.seed !== null ? Number(detailsObj.seed) : 42;
|
||||
shuffle(cocoImages, splitSeed);
|
||||
|
||||
// Split images
|
||||
const totalImages = cocoImages.length;
|
||||
const trainCount = Math.floor(totalImages * trainPercent / 100);
|
||||
const validCount = Math.floor(totalImages * validPercent / 100);
|
||||
const testCount = totalImages - trainCount - validCount;
|
||||
|
||||
const trainImages = cocoImages.slice(0, trainCount);
|
||||
const validImages = cocoImages.slice(trainCount, trainCount + validCount);
|
||||
const testImages = cocoImages.slice(trainCount + validCount);
|
||||
|
||||
// Helper to get image ids for each split
|
||||
const trainImageIds = new Set(trainImages.map(img => img.id));
|
||||
const validImageIds = new Set(validImages.map(img => img.id));
|
||||
const testImageIds = new Set(testImages.map(img => img.id));
|
||||
|
||||
// Split annotations
|
||||
const trainAnnotations = cocoAnnotations.filter(ann => trainImageIds.has(ann.image_id));
|
||||
const validAnnotations = cocoAnnotations.filter(ann => validImageIds.has(ann.image_id));
|
||||
const testAnnotations = cocoAnnotations.filter(ann => testImageIds.has(ann.image_id));
|
||||
|
||||
// Build final COCO JSONs with info section
|
||||
const buildCocoJson = (images, annotations, categories) => ({
|
||||
images,
|
||||
annotations,
|
||||
categories
|
||||
});
|
||||
|
||||
// Build COCO JSONs with info section
|
||||
const trainJson = buildCocoJson(trainImages, trainAnnotations, cocoCategories);
|
||||
const validJson = buildCocoJson(validImages, validAnnotations, cocoCategories);
|
||||
const testJson = buildCocoJson(testImages, testAnnotations, cocoCategories);
|
||||
|
||||
// Create output directory: projectname/trainingid/annotations
|
||||
const projectName = trainingProject && trainingProject.name ? trainingProject.name.replace(/\s+/g, '_') : `project_${detailsObj.project_id}`;
|
||||
const outDir = `${projectName}/${trainingId}`;
|
||||
const annotationsDir = `/home/kitraining/To_Annotate/annotations`;
|
||||
if (!fs.existsSync(annotationsDir)) {
|
||||
fs.mkdirSync(annotationsDir, { recursive: true });
|
||||
}
|
||||
|
||||
// Write to files in the annotations directory
|
||||
const trainPath = `${annotationsDir}/coco_project_${trainingId}_train.json`;
|
||||
const validPath = `${annotationsDir}/coco_project_${trainingId}_valid.json`;
|
||||
const testPath = `${annotationsDir}/coco_project_${trainingId}_test.json`;
|
||||
fs.writeFileSync(trainPath, JSON.stringify(trainJson, null, 2));
|
||||
fs.writeFileSync(validPath, JSON.stringify(validJson, null, 2));
|
||||
fs.writeFileSync(testPath, JSON.stringify(testJson, null, 2));
|
||||
console.log(`COCO JSON splits written to ${annotationsDir} for trainingId ${trainingId}`);
|
||||
|
||||
|
||||
|
||||
// Also generate inference exp.py in the same output directory as exp.py (project folder in workspace)
|
||||
const { generateYoloxInferenceExp } = require('./generate-yolox-exp');
|
||||
const path = require('path');
|
||||
const projectFolder = path.join(__dirname, '..', projectName, String(trainingId));
|
||||
if (!fs.existsSync(projectFolder)) {
|
||||
fs.mkdirSync(projectFolder, { recursive: true });
|
||||
}
|
||||
const inferenceExpPath = path.join(projectFolder, 'exp_infer.py');
|
||||
generateYoloxInferenceExp(trainingId).then(expContent => {
|
||||
fs.writeFileSync(inferenceExpPath, expContent);
|
||||
console.log(`Inference exp.py written to ${inferenceExpPath}`);
|
||||
}).catch(err => {
|
||||
console.error('Failed to generate inference exp.py:', err);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
module.exports = {generateTrainingJson};
|
||||
Reference in New Issue
Block a user