initial push

backend/services/generate-yolox-exp.js

@@ -0,0 +1,135 @@
+const fs = require('fs');
+const path = require('path');
+const Training = require('../models/training.js');
+const TrainingProject = require('../models/TrainingProject.js');
+
+// Remove Python comments and legacy code
+const exp_names = [
+    'YOLOX-s',
+    'YOLOX-m',
+    'YOLOX-l',
+    'YOLOX-x', 
+    'YOLOX-Darknet53', //todo
+    'YOLOX-Nano',
+    'YOLOX-Tiny'
+]
+
+//TODO: Clean up generation of exp_names.py and remove second exp creation!!!
+
+
+// Refactored: Accept trainingId, fetch info from DB
+async function generateYoloxExp(trainingId) {
+    // Fetch training row from DB by project_details_id if not found by PK
+    let training = await Training.findByPk(trainingId);
+    if (!training) {
+        training = await Training.findOne({ where: { project_details_id: trainingId } });
+    }
+    if (!training) throw new Error('Training not found for trainingId or project_details_id: ' + trainingId);
+
+    // If transfer_learning is 'coco', just return the path to the default exp.py
+    if (training.transfer_learning === 'coco') {
+        const selectedModel = training.selected_model.toLowerCase().replace('-', '_');
+        const expSourcePath = `/home/kitraining/Yolox/YOLOX-main/exps/default/${selectedModel}.py`;
+        if (!fs.existsSync(expSourcePath)) {
+            throw new Error(`Default exp.py not found for model: ${selectedModel} at ${expSourcePath}`);
+        }
+        // Copy to project folder (e.g., /home/kitraining/coco_tool/backend/project_XX/YY/exp.py)
+        const projectDetailsId = training.project_details_id;
+        const projectFolder = path.resolve(__dirname, `../project_23/${projectDetailsId}`);
+        if (!fs.existsSync(projectFolder)) {
+            fs.mkdirSync(projectFolder, { recursive: true });
+        }
+        const expDestPath = path.join(projectFolder, 'exp.py');
+        fs.copyFileSync(expSourcePath, expDestPath);
+        return { type: 'default', expPath: expDestPath };
+    }
+
+    // If transfer_learning is 'sketch', generate a custom exp.py as before
+    if (training.transfer_learning === 'sketch') {
+        // ...existing custom exp.py generation logic here (copy from previous implementation)...
+        // For brevity, you can call generateYoloxInferenceExp or similar here, or inline the logic.
+        // Example:
+        const expContent = await generateYoloxInferenceExp(trainingId);
+        return { type: 'custom', expContent };
+    }
+
+    throw new Error('Unknown transfer_learning type: ' + training.transfer_learning);
+}
+
+async function saveYoloxExp(trainingId, outPath) {
+    const expResult = await generateYoloxExp(trainingId);
+    if (expResult.type === 'custom' && expResult.expContent) {
+        fs.writeFileSync(outPath, expResult.expContent);
+        return outPath;
+    } else if (expResult.type === 'default' && expResult.expPath) {
+        // Optionally copy the file if outPath is different
+        if (expResult.expPath !== outPath) {
+            fs.copyFileSync(expResult.expPath, outPath);
+        }
+        return outPath;
+    } else {
+        throw new Error('Unknown expResult type or missing content');
+    }
+}
+
+// Generate a second exp.py for inference, using the provided template and DB values
+async function generateYoloxInferenceExp(trainingId, options = {}) {
+    let training = await Training.findByPk(trainingId);
+    if (!training) {
+        training = await Training.findOne({ where: { project_details_id: trainingId } });
+    }
+    if (!training) throw new Error('Training not found for trainingId or project_details_id: ' + trainingId);
+    // Always use the trainingId (project_details_id) for annotation file names
+    const projectDetailsId = training.project_details_id;
+    const dataDir = options.data_dir || '/home/kitraining/To_Annotate/';
+    const trainAnn = options.train_ann || `coco_project_${trainingId}_train.json`;
+    const valAnn = options.val_ann || `coco_project_${trainingId}_valid.json`;
+    const testAnn = options.test_ann || `coco_project_${trainingId}_test.json`;
+    // Get num_classes from TrainingProject.classes JSON
+    let numClasses = 80;
+    try {
+        const trainingProject = await TrainingProject.findByPk(projectDetailsId);
+        if (trainingProject && trainingProject.classes) {
+            let classesArr = trainingProject.classes;
+            if (typeof classesArr === 'string') {
+                classesArr = JSON.parse(classesArr);
+            }
+            if (Array.isArray(classesArr)) {
+                numClasses = classesArr.filter(c => c !== null && c !== undefined && c !== '').length;
+            } else if (typeof classesArr === 'object' && classesArr !== null) {
+                numClasses = Object.keys(classesArr).filter(k => classesArr[k] !== null && classesArr[k] !== undefined && classesArr[k] !== '').length;
+            }
+        }
+    } catch (e) {
+        console.warn('Could not determine num_classes from TrainingProject.classes:', e);
+    }
+    const depth = options.depth || training.depth || 1.00;
+    const width = options.width || training.width || 1.00;
+    const inputSize = options.input_size || training.input_size || [640, 640];
+    const mosaicScale = options.mosaic_scale || training.mosaic_scale || [0.1, 2];
+    const randomSize = options.random_size || training.random_size || [10, 20];
+    const testSize = options.test_size || training.test_size || [640, 640];
+    const expName = options.exp_name || 'inference_exp';
+    const enableMixup = options.enable_mixup !== undefined ? options.enable_mixup : false;
+    let expContent = '';
+    expContent += `#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n    def __init__(self):\n        super(Exp, self).__init__()\n        self.data_dir = "${dataDir}"\n        self.train_ann = "${trainAnn}"\n        self.val_ann = "${valAnn}"\n        self.test_ann = "coco_project_${trainingId}_test.json"\n        self.num_classes = ${numClasses}\n`;
+    // Set pretrained_ckpt if transfer_learning is 'coco'
+    if (training.transfer_learning && typeof training.transfer_learning === 'string' && training.transfer_learning.toLowerCase() === 'coco') {
+        const yoloxBaseDir = '/home/kitraining/Yolox/YOLOX-main';
+        const selectedModel = training.selected_model ? training.selected_model.replace(/\.pth$/i, '') : '';
+        if (selectedModel) {
+            expContent += `        self.pretrained_ckpt = r'${yoloxBaseDir}/pretrained/${selectedModel}.pth'\n`;
+        }
+    }
+    expContent += `        self.depth = ${depth}\n        self.width = ${width}\n        self.input_size = (${Array.isArray(inputSize) ? inputSize.join(', ') : inputSize})\n        self.mosaic_scale = (${Array.isArray(mosaicScale) ? mosaicScale.join(', ') : mosaicScale})\n        self.random_size = (${Array.isArray(randomSize) ? randomSize.join(', ') : randomSize})\n        self.test_size = (${Array.isArray(testSize) ? testSize.join(', ') : testSize})\n        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]\n        self.enable_mixup = ${enableMixup ? 'True' : 'False'}\n`;
+    return expContent;
+}
+
+// Save inference exp.py to a custom path
+async function saveYoloxInferenceExp(trainingId, outPath, options = {}) {
+    const expContent = await generateYoloxInferenceExp(trainingId, options);
+    fs.writeFileSync(outPath, expContent);
+    return outPath;
+}
+
+module.exports = { generateYoloxExp, saveYoloxExp, generateYoloxInferenceExp, saveYoloxInferenceExp };