feat: enable GPU selection for YOLO inference

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)

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) {

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"');
+  });
+});

views/training/yolovXInferenceSettings.ejs

@@ -160,6 +160,12 @@
                                                         <% } %>
                                                 </select>
                                             </div>
+                                            <div class="form-group col-md-2">
+                                                <h5><label for="device">Device:</label></h5>
+                                                <input type="text" id="device" name="device" value="1"
+                                                    style="width: 100%;">
+                                                <p>(cpu,gpu#,mps)</p>
+                                            </div>
                                         </div>
                                         <br>
                                         <div class="form-row">
@@ -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);