diff --git a/controllers/inference/datatovalues.py b/controllers/inference/datatovalues.py
index 4fd8c3ee..864c333c 100755
--- a/controllers/inference/datatovalues.py
+++ b/controllers/inference/datatovalues.py
@@ -235,19 +235,20 @@ def load_class_names(name_path):
darknet_path
+ f" {yolo_task}"
+ f" {yolo_mode}"
- + " predict model="
+ + ' predict model="'
+ weight_path
- + " source="
+ + '" source="'
+ image_path
- + " project="
+ + '" project="'
+ data_path
- + " name=output device="
+ + '" name=output device='
+ device
+ " save_txt=True save_conf=True"
+ " "
+ adv_options
- + " 2>&1 > "
+ + ' 2>&1 > "'
+ log_file
+ + '"'
)
print(cmd)
diff --git a/routes/inference/yoloInference.js b/routes/inference/yoloInference.js
index 409ffb39..2282e09f 100644
--- a/routes/inference/yoloInference.js
+++ b/routes/inference/yoloInference.js
@@ -3,7 +3,7 @@ const path = require("path");
async function yoloInference(req, res) {
try {
- const { exec } = require("child_process");
+ const { spawn } = require("child_process");
var date = Date.now();
@@ -14,7 +14,7 @@ async function yoloInference(req, res) {
yolovxPath = req.body.yolovx_path,
log = `${date}.log`,
inferenceFile = req.body.inference_file,
- device = req.body.device,
+ device = req.body.device || "cpu",
options = req.body.options,
yoloTask = req.body.yolo_task,
weightName = req.body.weights;
@@ -154,57 +154,43 @@ async function yoloInference(req, res) {
}
}
- var cmd = `python3 ${yoloScript} -d ${runPath} -i ${inferenceFilePath} -n ${classesPath} -l ${absUltralyticsProjectRun}/${log} -f ${ultralyticsPath} -w ${weightPath} -t ${yoloTask}`;
-
- var success = "";
- var error = "";
-
- fs.writeFileSync(`${absUltralyticsProjectRun}/${log}`, `${cmd}\n\n`);
+ const args = [
+ yoloScript,
+ "-d", runPath,
+ "-i", inferenceFilePath,
+ "-n", classesPath,
+ "-l", `${absUltralyticsProjectRun}/${log}`,
+ "-f", ultralyticsPath,
+ "-w", weightPath,
+ "-t", yoloTask,
+ "-D", device
+ ];
+
+ const child = spawn("python3", args);
+
+ child.stdout.on("data", (data) => {
+ global.logger.debug(`stdout: ${data}`);
+ });
- exec(cmd, (err, stdout, stderr) => {
- if (err) {
- global.logger.error(err);
- global.logger.debug(`This is the error: ${err.message}`);
-
- if (err.message != "stdout maxBuffer length exceeded") {
- success = err.message;
-
- fs.writeFile(
- `${ultralyticsProjectRun}/${errFile}`,
- success,
- (err) => {
- if (err) throw err;
- },
- );
- }
- } else if (stderr) {
- global.logger.debug(`This is the stderr: ${stderr}`);
-
- if (stderr != "stdout maxBuffer length exceeded") {
- fs.writeFile(
- `${ultralyticsProjectRun}/${errFile}`,
- stderr,
- (err) => {
- if (err) throw err;
- },
- );
- }
- }
+ child.stderr.on("data", (data) => {
+ global.logger.error(`stderr: ${data}`);
+ });
+ child.on("close", (code) => {
const completionData = {
- status: err ? 'error' : 'success',
+ status: code === 0 ? 'success' : 'error',
timestamp: date,
zipAvailable: fs.existsSync(`${runPath}/inference_results.zip`),
- csvAvailable: fs.existsSync(`${runPath}/inference_stats.csv`),
- message: success
+ csvAvailable: fs.existsSync(`${runPath}/inference_stats.csv`)
};
fs.writeFileSync(
`${runPath}/done.log`,
- `${cmd}\n\n${JSON.stringify(completionData, null, 2)}`
+ JSON.stringify(completionData, null, 2)
);
});
+ fs.writeFileSync(`${absUltralyticsProjectRun}/${log}`, `Arguments: ${args.join(" ")}\n\n`);
res.send({ Success: `YOLO Inference Started` });
} catch (err) {
diff --git a/tests/integration/yoloInference.test.js b/tests/integration/yoloInference.test.js
new file mode 100644
index 00000000..aec7a56a
--- /dev/null
+++ b/tests/integration/yoloInference.test.js
@@ -0,0 +1,169 @@
+const request = require('supertest');
+
+// Set up mocks before requiring app
+jest.mock('decompress-zip', () => jest.fn());
+jest.mock('decompress-zip/lib/extractors', () => ({
+ folder: jest.fn(),
+}));
+jest.mock('ffmpeg', () => jest.fn());
+jest.mock('sharp', () => jest.fn());
+jest.mock('unzipper', () => jest.fn());
+
+const mockExec = jest.fn((cmd, cb) => cb(null, '', ''));
+jest.mock('child_process', () => ({
+ exec: mockExec,
+}));
+
+jest.mock('sqlite3', () => {
+ const mockDbInstance = {
+ run: jest.fn((sql, params, cb) => {
+ const callback = typeof params === 'function' ? params : cb;
+ if (typeof callback === 'function') callback(null);
+ }),
+ get: jest.fn((sql, params, cb) => {
+ const callback = typeof params === 'function' ? params : cb;
+ if (typeof callback === 'function') callback(null, {});
+ }),
+ all: jest.fn((sql, params, cb) => {
+ const callback = typeof params === 'function' ? params : cb;
+ if (typeof callback === 'function') callback(null, []);
+ }),
+ close: jest.fn((cb) => {
+ if (typeof cb === 'function') cb(null);
+ }),
+ };
+ const sqlite3Mock = {
+ OPEN_READWRITE: 1,
+ OPEN_CREATE: 2,
+ Database: jest.fn(() => mockDbInstance),
+ verbose: () => sqlite3Mock,
+ };
+ return sqlite3Mock;
+});
+
+jest.mock('socket.io-client', () => ({
+ protocol: 'http',
+}));
+
+// Mock fs functions
+jest.mock('fs', () => ({
+ existsSync: jest.fn().mockReturnValue(true),
+ mkdirSync: jest.fn(),
+ writeFileSync: jest.fn(),
+ copyFileSync: jest.fn(),
+ symlinkSync: jest.fn(),
+ readFileSync: jest.fn().mockReturnValue(Buffer.from('mockImage')),
+ writeFile: jest.fn((path, data, callback) => callback(null)),
+}));
+
+// Mock probe module
+jest.mock('probe-image-size', () => ({
+ sync: jest.fn(() => ({ width: 800, height: 600 })),
+}));
+
+// Mock queries
+const mockQueries = {
+ project: {
+ getAllClasses: jest.fn().mockResolvedValue({ rows: [{ CName: 'class1' }] }),
+ getAllImages: jest.fn().mockResolvedValue({ rows: [{ IName: 'image1.jpg' }] }),
+ getLabelsForImageName: jest.fn().mockResolvedValue({
+ rows: [{ CName: 'class1', X: 10, Y: 10, W: 100, H: 100 }]
+ }),
+ },
+};
+jest.mock('../../queries/queries', () => mockQueries);
+
+// Define global variables expected by the application/routes
+global.fs = require('fs');
+global.probe = require('probe-image-size');
+global.currentPath = '/mock/current/path/';
+global.configFile = {};
+
+const app = require('../../app');
+
+describe('YOLO Inference Integration Tests', () => {
+ beforeEach(() => {
+ jest.clearAllMocks();
+ });
+
+ it('should retrieve custom device parameter and pass it to python command with -D', async () => {
+ const payload = {
+ PName: 'test-project',
+ Admin: 'admin-user',
+ yolovx_path: '/path/to/yolovx',
+ inference_file: 'inference_image.jpg',
+ device: 'gpu',
+ options: '',
+ yolo_task: 'detect',
+ weights: 'best.pt'
+ };
+
+ const response = await request(app)
+ .post('/yolo-inf')
+ .set('Cookie', ['Username=test-user'])
+ .send(payload);
+
+ expect(response.status).toBe(200);
+ expect(response.body).toEqual({ Success: 'YOLO Inference Started' });
+ expect(mockExec).toHaveBeenCalled();
+
+ // Retrieve the executed command argument
+ const executedCmd = mockExec.mock.calls[0][0];
+ expect(executedCmd).toContain('-D gpu');
+ });
+
+ it('should default device parameter to cpu when not provided and pass it to python command with -D', async () => {
+ const payload = {
+ PName: 'test-project',
+ Admin: 'admin-user',
+ yolovx_path: '/path/to/yolovx',
+ inference_file: 'inference_image.jpg',
+ options: '',
+ yolo_task: 'detect',
+ weights: 'best.pt'
+ };
+
+ const response = await request(app)
+ .post('/yolo-inf')
+ .set('Cookie', ['Username=test-user'])
+ .send(payload);
+
+ expect(response.status).toBe(200);
+ expect(response.body).toEqual({ Success: 'YOLO Inference Started' });
+ expect(mockExec).toHaveBeenCalled();
+
+ // Retrieve the executed command argument
+ const executedCmd = mockExec.mock.calls[0][0];
+ expect(executedCmd).toContain('-D cpu');
+ });
+
+ it('should wrap path parameters with double quotes when executing the command, handling weights with parentheses', async () => {
+ const payload = {
+ PName: 'test-project',
+ Admin: 'admin-user',
+ yolovx_path: '/path/to/yolovx path',
+ inference_file: 'inference_image(1).jpg',
+ device: 'gpu',
+ options: '',
+ yolo_task: 'detect',
+ weights: 'best(1).pt'
+ };
+
+ const response = await request(app)
+ .post('/yolo-inf')
+ .set('Cookie', ['Username=test-user'])
+ .send(payload);
+
+ expect(response.status).toBe(200);
+ expect(response.body).toEqual({ Success: 'YOLO Inference Started' });
+ expect(mockExec).toHaveBeenCalled();
+
+ // Retrieve the executed command argument
+ const executedCmd = mockExec.mock.calls[0][0];
+
+ // Check that paths are wrapped in double quotes
+ expect(executedCmd).toContain('-w "/mock/current/path/public/projects/admin-user-test-project/training/weights/best(1).pt"');
+ expect(executedCmd).toContain('-i "inference_image(1).jpg"');
+ expect(executedCmd).toContain('-f "/path/to/yolovx path"');
+ });
+});
diff --git a/views/training/yolovXInferenceSettings.ejs b/views/training/yolovXInferenceSettings.ejs
index 7d1ce7f4..4e3e28a8 100644
--- a/views/training/yolovXInferenceSettings.ejs
+++ b/views/training/yolovXInferenceSettings.ejs
@@ -160,6 +160,12 @@
<% } %>
+
@@ -344,7 +350,7 @@
formData.append("yolo_task", document.getElementById("yolo_task").value);
formData.append("yolo_mode", document.getElementById("yolo_mode").value);
// formData.append("options", document.getElementById("options").value);
- // formData.append("device", document.getElementById("device").value);
+ formData.append("device", document.getElementById("device").value);
console.log(formData);