497 lines
21 KiB
JavaScript
497 lines
21 KiB
JavaScript
const express = require('express');
|
|
const multer = require('multer');
|
|
const upload = multer();
|
|
const TrainingProject = require('../models/TrainingProject.js');
|
|
const LabelStudioProject = require('../models/LabelStudioProject.js')
|
|
const { seedLabelStudio, updateStatus } = require('../services/seed-label-studio.js');
|
|
const fs = require('fs');
|
|
const path = require('path');
|
|
const {generateTrainingJson} = require('../services/generate-json-yolox.js')
|
|
|
|
|
|
const router = express.Router();
|
|
|
|
// Ensure JSON bodies are parsed for all routes
|
|
router.use(express.json());
|
|
|
|
router.get('/seed', async (req, res) => {
|
|
const result = await seedLabelStudio();
|
|
res.json(result);
|
|
});
|
|
|
|
|
|
|
|
// Trigger generate-json-yolox.js
|
|
|
|
router.post('/generate-yolox-json', async (req, res) => {
|
|
const { project_id } = req.body;
|
|
if (!project_id) {
|
|
return res.status(400).json({ message: 'Missing project_id in request body' });
|
|
}
|
|
try {
|
|
// Generate COCO JSONs
|
|
// Find all TrainingProjectDetails for this project
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const detailsRows = await TrainingProjectDetails.findAll({ where: { project_id } });
|
|
if (!detailsRows || detailsRows.length === 0) {
|
|
return res.status(404).json({ message: 'No TrainingProjectDetails found for project ' + project_id });
|
|
}
|
|
// For each details row, generate coco.jsons and exp.py in projectfolder/project_details_id
|
|
const Training = require('../models/training.js');
|
|
const { saveYoloxExp } = require('../services/generate-yolox-exp.js');
|
|
const TrainingProject = require('../models/TrainingProject.js');
|
|
const trainingProject = await TrainingProject.findByPk(project_id);
|
|
const projectName = trainingProject.name ? trainingProject.name.replace(/\s+/g, '_') : `project_${project_id}`;
|
|
for (const details of detailsRows) {
|
|
const detailsId = details.id;
|
|
await generateTrainingJson(detailsId);
|
|
const trainings = await Training.findAll({ where: { project_details_id: detailsId } });
|
|
if (trainings.length === 0) continue;
|
|
// For each training, save exp.py in projectfolder/project_details_id
|
|
const outDir = path.join(__dirname, '..', projectName, String(detailsId));
|
|
if (!fs.existsSync(outDir)) fs.mkdirSync(outDir, { recursive: true });
|
|
for (const training of trainings) {
|
|
const expFilePath = path.join(outDir, 'exp.py');
|
|
await saveYoloxExp(training.id, expFilePath);
|
|
}
|
|
}
|
|
|
|
// Find all trainings for this project
|
|
// ...existing code...
|
|
res.json({ message: 'YOLOX JSON and exp.py generated for project ' + project_id });
|
|
} catch (err) {
|
|
console.error('Error generating YOLOX JSON:', err);
|
|
res.status(500).json({ message: 'Failed to generate YOLOX JSON', error: err.message });
|
|
}
|
|
});
|
|
|
|
// Start YOLOX training
|
|
const { spawn } = require('child_process');
|
|
router.post('/start-yolox-training', async (req, res) => {
|
|
try {
|
|
const { project_id, training_id } = req.body;
|
|
// Get project name
|
|
const trainingProject = await TrainingProject.findByPk(project_id);
|
|
const projectName = trainingProject.name ? trainingProject.name.replace(/\s+/g, '_') : `project_${project_id}`;
|
|
// Look up training row by id or project_details_id
|
|
const Training = require('../models/training.js');
|
|
let trainingRow = await Training.findByPk(training_id);
|
|
if (!trainingRow) {
|
|
trainingRow = await Training.findOne({ where: { project_details_id: training_id } });
|
|
}
|
|
if (!trainingRow) {
|
|
return res.status(404).json({ error: `Training row not found for id or project_details_id ${training_id}` });
|
|
}
|
|
const project_details_id = trainingRow.project_details_id;
|
|
// Use the generated exp.py from the correct project folder
|
|
const outDir = path.join(__dirname, '..', projectName, String(project_details_id));
|
|
const yoloxMainDir = '/home/kitraining/Yolox/YOLOX-main';
|
|
const expSrc = path.join(outDir, 'exp.py');
|
|
if (!fs.existsSync(expSrc)) {
|
|
return res.status(500).json({ error: `exp.py not found at ${expSrc}` });
|
|
}
|
|
// Source venv and run YOLOX training in YOLOX-main folder
|
|
const yoloxVenv = '/home/kitraining/Yolox/yolox_venv/bin/activate';
|
|
// Determine model argument based on selected_model and transfer_learning
|
|
let modelArg = '';
|
|
let cmd = '';
|
|
if (
|
|
trainingRow.transfer_learning &&
|
|
typeof trainingRow.transfer_learning === 'string' &&
|
|
trainingRow.transfer_learning.toLowerCase() === 'coco'
|
|
) {
|
|
// If transfer_learning is 'coco', add -o and modelArg
|
|
modelArg = ` -c /home/kitraining/Yolox/YOLOX-main/pretrained/${trainingRow.selected_model}`;
|
|
cmd = `bash -c 'source ${yoloxVenv} && python tools/train.py -f ${expSrc} -d 1 -b 8 --fp16 -o ${modelArg}.pth --cache'`;
|
|
} else if (
|
|
trainingRow.selected_model &&
|
|
trainingRow.selected_model.toLowerCase() === 'coco' &&
|
|
(!trainingRow.transfer_learning || trainingRow.transfer_learning === false)
|
|
) {
|
|
// If selected_model is 'coco' and not transfer_learning, add modelArg only
|
|
modelArg = ` -c /pretrained/${trainingRow.selected_model}`;
|
|
cmd = `bash -c 'source ${yoloxVenv} && python tools/train.py -f ${expSrc} -d 1 -b 8 --fp16 -o ${modelArg}.pth --cache'`;
|
|
} else {
|
|
// Default: no modelArg
|
|
cmd = `bash -c 'source ${yoloxVenv} && python tools/train.py -f ${expSrc} -d 1 -b 8 --fp16' --cache`;
|
|
}
|
|
console.log(cmd)
|
|
const child = spawn(cmd, { shell: true, cwd: yoloxMainDir });
|
|
child.stdout.pipe(process.stdout);
|
|
child.stderr.pipe(process.stderr);
|
|
|
|
res.json({ message: 'Training started' });
|
|
} catch (err) {
|
|
res.status(500).json({ error: 'Failed to start training', details: err.message });
|
|
}
|
|
});
|
|
|
|
// Get YOLOX training log
|
|
router.get('/training-log', async (req, res) => {
|
|
try {
|
|
const { project_id, training_id } = req.query;
|
|
const trainingProject = await TrainingProject.findByPk(project_id);
|
|
const projectName = trainingProject.name ? trainingProject.name.replace(/\s+/g, '_') : `project_${project_id}`;
|
|
const outDir = path.join(__dirname, '..', projectName, String(training_id));
|
|
const logPath = path.join(outDir, 'training.log');
|
|
if (!fs.existsSync(logPath)) {
|
|
return res.status(404).json({ error: 'Log not found' });
|
|
}
|
|
const logData = fs.readFileSync(logPath, 'utf8');
|
|
res.json({ log: logData });
|
|
} catch (err) {
|
|
res.status(500).json({ error: 'Failed to fetch log', details: err.message });
|
|
}
|
|
});
|
|
|
|
router.post('/training-projects', upload.single('project_image'), async (req, res) => {
|
|
try {
|
|
const { title, description } = req.body;
|
|
const classes = JSON.parse(req.body.classes);
|
|
const project_image = req.file ? req.file.buffer : null;
|
|
const project_image_type = req.file ? req.file.mimetype : null;
|
|
await TrainingProject.create({
|
|
title,
|
|
description,
|
|
classes,
|
|
project_image,
|
|
project_image_type
|
|
});
|
|
res.json({ message: 'Project created!' });
|
|
} catch (error) {
|
|
console.error('Error creating project:', error);
|
|
res.status(500).json({ message: 'Failed to create project', error: error.message });
|
|
}
|
|
});
|
|
|
|
router.get('/training-projects', async (req, res) => {
|
|
try {
|
|
const projects = await TrainingProject.findAll();
|
|
// Convert BLOB to base64 data URL for each project
|
|
const serialized = projects.map(project => {
|
|
const plain = project.get({ plain: true });
|
|
if (plain.project_image) {
|
|
const base64 = Buffer.from(plain.project_image).toString('base64');
|
|
const mimeType = plain.project_image_type || 'image/png';
|
|
plain.project_image = `data:${mimeType};base64,${base64}`;
|
|
}
|
|
return plain;
|
|
});
|
|
res.json(serialized);
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to fetch projects', error: error.message });
|
|
}
|
|
});
|
|
|
|
router.get('/update-status', async (req, res) => {
|
|
res.json(updateStatus)
|
|
})
|
|
|
|
router.get('/label-studio-projects', async (req, res) => {
|
|
try {
|
|
const LabelStudioProject = require('../models/LabelStudioProject.js');
|
|
const Image = require('../models/Images.js');
|
|
const Annotation = require('../models/Annotation.js');
|
|
const labelStudioProjects = await LabelStudioProject.findAll();
|
|
const projectsWithCounts = await Promise.all(labelStudioProjects.map(async project => {
|
|
const plain = project.get({ plain: true });
|
|
// Get all images for this project
|
|
const images = await Image.findAll({ where: { project_id: plain.project_id } });
|
|
let annotationCounts = {};
|
|
if (images.length > 0) {
|
|
const imageIds = images.map(img => img.image_id);
|
|
// Get all annotations for these images
|
|
const annotations = await Annotation.findAll({ where: { image_id: imageIds } });
|
|
// Count by label
|
|
for (const ann of annotations) {
|
|
const label = ann.Label;
|
|
annotationCounts[label] = (annotationCounts[label] || 0) + 1;
|
|
}
|
|
}
|
|
plain.annotationCounts = annotationCounts;
|
|
return plain;
|
|
}));
|
|
res.json(projectsWithCounts);
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to fetch projects', error: error.message });
|
|
}
|
|
})
|
|
|
|
|
|
// POST endpoint to create TrainingProjectDetails with all fields
|
|
router.post('/training-project-details', async (req, res) => {
|
|
try {
|
|
const {
|
|
project_id,
|
|
annotation_projects,
|
|
class_map,
|
|
description
|
|
} = req.body;
|
|
if (!project_id || !annotation_projects) {
|
|
return res.status(400).json({ message: 'Missing required fields' });
|
|
}
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const created = await TrainingProjectDetails.create({
|
|
project_id,
|
|
annotation_projects,
|
|
class_map: class_map || null,
|
|
description: description || null
|
|
});
|
|
res.json({ message: 'TrainingProjectDetails created', details: created });
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to create TrainingProjectDetails', error: error.message });
|
|
}
|
|
});
|
|
|
|
// GET endpoint to fetch all TrainingProjectDetails
|
|
router.get('/training-project-details', async (req, res) => {
|
|
try {
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const details = await TrainingProjectDetails.findAll();
|
|
res.json(details);
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to fetch TrainingProjectDetails', error: error.message });
|
|
}
|
|
});
|
|
|
|
// PUT endpoint to update class_map and description in TrainingProjectDetails
|
|
router.put('/training-project-details', async (req, res) => {
|
|
try {
|
|
const { project_id, class_map, description } = req.body;
|
|
if (!project_id || !class_map || !description) {
|
|
return res.status(400).json({ message: 'Missing required fields' });
|
|
}
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const details = await TrainingProjectDetails.findOne({ where: { project_id } });
|
|
if (!details) {
|
|
return res.status(404).json({ message: 'TrainingProjectDetails not found' });
|
|
}
|
|
details.class_map = class_map;
|
|
details.description = description;
|
|
await details.save();
|
|
res.json({ message: 'Class map and description updated', details });
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to update class map or description', error: error.message });
|
|
}
|
|
});
|
|
|
|
// POST endpoint to receive YOLOX settings and save to DB (handles multipart/form-data)
|
|
router.post('/yolox-settings', upload.any(), async (req, res) => {
|
|
try {
|
|
const settings = req.body;
|
|
// Debug: Log all received fields and types
|
|
console.log('--- YOLOX settings received ---');
|
|
console.log('settings:', settings);
|
|
if (req.files && req.files.length > 0) {
|
|
console.log('Files received:', req.files.map(f => ({ fieldname: f.fieldname, originalname: f.originalname, size: f.size })));
|
|
}
|
|
// Declare requiredFields once
|
|
const requiredFields = ['project_details_id', 'exp_name', 'max_epoch', 'depth', 'width', 'activation', 'train', 'valid', 'test', 'selected_model', 'transfer_learning'];
|
|
// Log types of required fields
|
|
requiredFields.forEach(field => {
|
|
console.log(`Field '${field}': value='${settings[field]}', type='${typeof settings[field]}'`);
|
|
});
|
|
// Map select_model to selected_model if present
|
|
if (settings && settings.select_model && !settings.selected_model) {
|
|
settings.selected_model = settings.select_model;
|
|
delete settings.select_model;
|
|
}
|
|
// Lookup project_details_id from project_id
|
|
if (!settings.project_id || isNaN(Number(settings.project_id))) {
|
|
throw new Error('Missing or invalid project_id in request. Cannot assign training to a project.');
|
|
}
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
let details = await TrainingProjectDetails.findOne({ where: { project_id: settings.project_id } });
|
|
if (!details) {
|
|
details = await TrainingProjectDetails.create({
|
|
project_id: settings.project_id,
|
|
annotation_projects: [],
|
|
class_map: null,
|
|
description: null
|
|
});
|
|
}
|
|
settings.project_details_id = details.id;
|
|
// Map 'act' from frontend to 'activation' for DB
|
|
if (settings.act !== undefined) {
|
|
settings.activation = settings.act;
|
|
delete settings.act;
|
|
}
|
|
// Type conversion for DB compatibility
|
|
[
|
|
'max_epoch', 'depth', 'width', 'warmup_epochs', 'warmup_lr', 'no_aug_epochs', 'min_lr_ratio', 'weight_decay', 'momentum', 'print_interval', 'eval_interval', 'test_conf', 'nmsthre', 'multiscale_range', 'degrees', 'translate', 'shear', 'train', 'valid', 'test'
|
|
].forEach(f => {
|
|
if (settings[f] !== undefined) settings[f] = Number(settings[f]);
|
|
});
|
|
// Improved boolean conversion
|
|
['ema', 'enable_mixup', 'save_history_ckpt'].forEach(f => {
|
|
if (settings[f] !== undefined) {
|
|
if (typeof settings[f] === 'string') {
|
|
settings[f] = settings[f].toLowerCase() === 'true';
|
|
} else {
|
|
settings[f] = Boolean(settings[f]);
|
|
}
|
|
}
|
|
});
|
|
// Improved array conversion
|
|
['mosaic_scale', 'mixup_scale', 'scale'].forEach(f => {
|
|
if (settings[f] && typeof settings[f] === 'string') {
|
|
settings[f] = settings[f]
|
|
.split(',')
|
|
.map(s => Number(s.trim()))
|
|
.filter(n => !isNaN(n));
|
|
}
|
|
});
|
|
// Trim all string fields
|
|
Object.keys(settings).forEach(f => {
|
|
if (typeof settings[f] === 'string') settings[f] = settings[f].trim();
|
|
});
|
|
// Set default for transfer_learning if missing
|
|
if (settings.transfer_learning === undefined) settings.transfer_learning = false;
|
|
// Convert empty string seed to null
|
|
if ('seed' in settings && (settings.seed === '' || settings.seed === undefined)) {
|
|
settings.seed = null;
|
|
}
|
|
// Validate required fields for training table
|
|
for (const field of requiredFields) {
|
|
if (settings[field] === undefined || settings[field] === null || settings[field] === '') {
|
|
console.error('Missing required field:', field, 'Value:', settings[field]);
|
|
throw new Error('Missing required field: ' + field);
|
|
}
|
|
}
|
|
console.log('Received YOLOX settings:', settings);
|
|
// Handle uploaded model file (ckpt_upload)
|
|
if (req.files && req.files.length > 0) {
|
|
const ckptFile = req.files.find(f => f.fieldname === 'ckpt_upload');
|
|
if (ckptFile) {
|
|
const uploadDir = path.join(__dirname, '..', 'uploads');
|
|
if (!fs.existsSync(uploadDir)) fs.mkdirSync(uploadDir);
|
|
const filename = ckptFile.originalname || `uploaded_model_${settings.project_id}.pth`;
|
|
const filePath = path.join(uploadDir, filename);
|
|
fs.writeFileSync(filePath, ckptFile.buffer);
|
|
settings.model_upload = filePath;
|
|
}
|
|
}
|
|
// Save settings to DB only (no file)
|
|
const { pushYoloxExpToDb } = require('../services/push-yolox-exp.js');
|
|
const training = await pushYoloxExpToDb(settings);
|
|
res.json({ message: 'YOLOX settings saved to DB', training });
|
|
} catch (error) {
|
|
console.error('Error in /api/yolox-settings:', error.stack || error);
|
|
res.status(500).json({ message: 'Failed to save YOLOX settings', error: error.message });
|
|
}
|
|
});
|
|
|
|
// POST endpoint to receive binary model file and save to disk (not DB)
|
|
router.post('/yolox-settings/upload', async (req, res) => {
|
|
try {
|
|
const projectId = req.query.project_id;
|
|
if (!projectId) return res.status(400).json({ message: 'Missing project_id in query' });
|
|
// Save file to disk
|
|
const uploadDir = path.join(__dirname, '..', 'uploads');
|
|
if (!fs.existsSync(uploadDir)) fs.mkdirSync(uploadDir);
|
|
const filename = req.headers['x-upload-filename'] || `uploaded_model_${projectId}.pth`;
|
|
const filePath = path.join(uploadDir, filename);
|
|
const chunks = [];
|
|
req.on('data', chunk => chunks.push(chunk));
|
|
req.on('end', async () => {
|
|
const buffer = Buffer.concat(chunks);
|
|
fs.writeFile(filePath, buffer, async err => {
|
|
if (err) {
|
|
console.error('Error saving file:', err);
|
|
return res.status(500).json({ message: 'Failed to save model file', error: err.message });
|
|
}
|
|
// Update latest training row for this project with file path
|
|
try {
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const Training = require('../models/training.js');
|
|
// Find details row for this project
|
|
const details = await TrainingProjectDetails.findOne({ where: { project_id: projectId } });
|
|
if (!details) return res.status(404).json({ message: 'No TrainingProjectDetails found for project_id' });
|
|
// Find latest training for this details row
|
|
const training = await Training.findOne({ where: { project_details_id: details.id }, order: [['createdAt', 'DESC']] });
|
|
if (!training) return res.status(404).json({ message: 'No training found for project_id' });
|
|
// Save file path to model_upload field
|
|
training.model_upload = filePath;
|
|
await training.save();
|
|
res.json({ message: 'Model file uploaded and saved to disk', filename, trainingId: training.id });
|
|
} catch (dbErr) {
|
|
console.error('Error updating training with file path:', dbErr);
|
|
res.status(500).json({ message: 'File saved but failed to update training row', error: dbErr.message });
|
|
}
|
|
});
|
|
});
|
|
} catch (error) {
|
|
console.error('Error in /api/yolox-settings/upload:', error.stack || error);
|
|
res.status(500).json({ message: 'Failed to upload model file', error: error.message });
|
|
}
|
|
});
|
|
|
|
// GET endpoint to fetch all trainings (optionally filtered by project_id)
|
|
router.get('/trainings', async (req, res) => {
|
|
try {
|
|
const project_id = req.query.project_id;
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const Training = require('../models/training.js');
|
|
if (project_id) {
|
|
// Find all details rows for this project
|
|
const detailsRows = await TrainingProjectDetails.findAll({ where: { project_id } });
|
|
if (!detailsRows || detailsRows.length === 0) return res.json([]);
|
|
// Get all trainings linked to any details row for this project
|
|
const detailsIds = detailsRows.map(d => d.id);
|
|
const trainings = await Training.findAll({ where: { project_details_id: detailsIds } });
|
|
return res.json(trainings);
|
|
} else {
|
|
// Return all trainings if no project_id is specified
|
|
const trainings = await Training.findAll();
|
|
return res.json(trainings);
|
|
}
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to fetch trainings', error: error.message });
|
|
}
|
|
});
|
|
|
|
// DELETE endpoint to remove a training by id
|
|
router.delete('/trainings/:id', async (req, res) => {
|
|
try {
|
|
const Training = require('../models/training.js');
|
|
const id = req.params.id;
|
|
const deleted = await Training.destroy({ where: { id } });
|
|
if (deleted) {
|
|
res.json({ message: 'Training deleted' });
|
|
} else {
|
|
res.status(404).json({ message: 'Training not found' });
|
|
}
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to delete training', error: error.message });
|
|
}
|
|
});
|
|
|
|
// DELETE endpoint to remove a training project and all related entries
|
|
router.delete('/training-projects/:id', async (req, res) => {
|
|
try {
|
|
const projectId = req.params.id;
|
|
const TrainingProject = require('../models/TrainingProject.js');
|
|
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js');
|
|
const Training = require('../models/training.js');
|
|
// Find details row(s) for this project
|
|
const detailsRows = await TrainingProjectDetails.findAll({ where: { project_id: projectId } });
|
|
const detailsIds = detailsRows.map(d => d.id);
|
|
// Delete all trainings linked to these details
|
|
if (detailsIds.length > 0) {
|
|
await Training.destroy({ where: { project_details_id: detailsIds } });
|
|
await TrainingProjectDetails.destroy({ where: { project_id: projectId } });
|
|
}
|
|
// Delete the project itself
|
|
const deleted = await TrainingProject.destroy({ where: { project_id: projectId } });
|
|
if (deleted) {
|
|
res.json({ message: 'Training project and all related entries deleted' });
|
|
} else {
|
|
res.status(404).json({ message: 'Training project not found' });
|
|
}
|
|
} catch (error) {
|
|
res.status(500).json({ message: 'Failed to delete training project', error: error.message });
|
|
}
|
|
});
|
|
|
|
module.exports = router;
|