Run resection identification using Docker container

toolbox/core/bst_containers.m

@@ -5,7 +5,7 @@
 %  [errMsg, engineName]    = bst_containers('GetEngine')
 %  [errMsg, imageList]     = bst_containers('GetImages')
 %  [errMsg, imageSha]      = bst_containers('ImportImage', imageSource, [imageTag])
-%   errMsg                 = bst_containers('RunContainer', containerName, imageSha, [volumes], [isDaemon])
+%   errMsg                 = bst_containers('RunContainer', containerName, imageSha, [volumes], [isDaemon], [containerArgs])
 %  [errMsg, cmdout]        = bst_containers('ExecInContainer', containerName, cmdStr)
 %  [errMsg, containerInfo] = bst_containers('GetContainerInfo', containerName)
 %   errMsg                 = bst_containers('StopContainer', containerName, [isForced=0])
@@ -248,9 +248,12 @@
 
 
 %% ===== RUN CONTAINER AS DAEMON =====
-function errMsg = RunContainer(containerName, imageSha, volumes, isDaemon)
-% USAGE: errMsg = bst_containers('RunContainer', containerName, imageSha, volumes, isDaemon)
+function errMsg = RunContainer(containerName, imageSha, volumes, isDaemon, containerArgs)
+% USAGE: errMsg = bst_containers('RunContainer', containerName, imageSha, volumes, isDaemon, containerArgs)
     % Validate inputs
+    if nargin < 5 || isempty(containerArgs)
+        containerArgs = '';
+    end
     if nargin < 4 || isempty(isDaemon)
         isDaemon = 0;
     end
@@ -275,15 +278,21 @@
         volumesStr = strjoin(pairs, ' ');
     end
 
+    % Use GPU with container engine
+    gpuStr = '';
+    if bst_get('ContainerUseGpu') && system('which nvidia-smi') == 0
+        gpuStr = '--gpus all';
+    end
+
     % Run container
     switch engineName
         case 'docker'
             if ~isDaemon
                 % Run ENTRYPOINT
-                cmdStr = sprintf('docker run --rm --name %s %s %s', containerName, volumesStr, imageSha);
+                cmdStr = sprintf('docker run --rm --name %s %s %s %s %s', containerName, gpuStr, volumesStr, imageSha, containerArgs);
             else
                 % Replace ENTRYPOINT (if any) with `sleep infinity`
-                cmdStr = sprintf('docker run -d --name %s %s --entrypoint sleep %s infinity', containerName, volumesStr, imageSha);
+                cmdStr = sprintf('docker run -d --name %s %s %s --entrypoint sleep %s infinity', containerName, gpuStr, volumesStr, imageSha);
             end
             [status, cmdout] = system(cmdStr);
     end
@@ -308,6 +317,10 @@
         return
     end
 
+    % Flag to track interruption
+    processState = containers.Map({'isInterruptCleanup'}, {1});
+    % Clean up on function end, errors or Ctrl+C is pressed
+    cleanupObj = onCleanup(@() ProcessInterrupted(containerName, processState));
     % Run command
     switch engineName
         case 'docker'
@@ -316,11 +329,15 @@
             else
                 commandWrapper = ''''; % Single quote
             end
-            [status, cmdout] = system(['docker exec ' containerName ' sh -c ' commandWrapper cmdStr commandWrapper]);
+            % Execute the running container
+            commandExec = ['docker exec ' containerName ' sh -c ' commandWrapper cmdStr commandWrapper];
+            [status, cmdout] = system(commandExec, '-echo');
             if status ~= 0
                 errMsg = strtrim(cmdout);
             end
     end
+    % Code in container ended normally
+    processState('isInterruptCleanup') = 0;
 end
 
 
@@ -428,6 +445,15 @@
 end
 
 
+%% ===== PROCESS INTERRUPTED =====
+function ProcessInterrupted(containerName, processState)
+    if processState('isInterruptCleanup')
+        bst_plugin('Unload', regexprep(containerName, '^bst_', ''));
+        bst_error('The process running in the container was interrupted', 'Container', 0);
+    end
+end
+
+
 %% ===== GET ONLINE MANIFEST DIGEST =====
 function [errMsg, manifestSha] = GetOnlineManifest(imageSource)
     manifestSha = '';

toolbox/core/bst_plugin.m

@@ -209,15 +209,9 @@
     PlugDesc(end+1)              = GetStruct('resection-identification');
     PlugDesc(end).Version        = 'latest';
     PlugDesc(end).Category       = 'Anatomy';
-    PlugDesc(end).AutoUpdate     = 1;
-    PlugDesc(end).URLzip         = ['https://neuroimage.usc.edu/bst/getupdate.php?d=bst_resection_identification_' OsType '.zip'];
-    PlugDesc(end).TestFile       = 'resection_identification';
-    if strcmp(OsType, 'win64')
-        PlugDesc(end).TestFile   = [PlugDesc(end).TestFile, '.bat'];
-    end
-    PlugDesc(end).URLinfo        = 'https://github.com/ajoshiusc/auto_resection_mask/tree/brainstorm-plugin';
+    PlugDesc(end).URLinfo        = 'https://github.com/ajoshiusc/auto_resection_mask/tree/brainstorm-container';
+    PlugDesc(end).ImageSource    = ['docker.io/brainstormtools/auto-resection-mask:' PlugDesc(end).Version];
     PlugDesc(end).CompiledStatus = 1;
-    PlugDesc(end).LoadFolders    = {'bin'};
 
     % === ANATOMY: ROAST ===
     PlugDesc(end+1)              = GetStruct('roast');

toolbox/process/functions/process_resection_identification.m

@@ -130,39 +130,42 @@
         errMsg = 'The fiducials (NAS, LPA, RPA) are missing in the pre-op (default) MRI. Set them first before proceeding.';
         return;
     end
-
-    % Install/load resection-identification plugin
-    [isOk, errInstall, PlugDesc] = bst_plugin('Install', 'resection-identification', isInteractive);
-    if ~isOk
-        errMsg = [errMsg, errInstall];
-        return;
-    end
-
+
+    % === CALL RESECTION-IDENTIFICATION PIPELINE ===
+    % Container plugin name
+    plugName = 'resection-identification';
+    bst_progress('text', ['Calling ' plugName]);
+    tic;    
+    % Ensure container plugin: Installs and/or Loads
+    % Install container plugin === Import image into container engine
+    % Load    container plugin === Run container in standby (name ['bst_' plugName])
+    ensureRes = bst_plugin('Ensure', plugName);
+    % Retrieve info of container
+    [errMsg, containerInfo] = bst_containers('GetContainerInfo', ['bst_' plugName]);
+
     % === SAVE BOTH MRI AS NIfTI ===
-    bst_progress('text', 'Exporting pre- and post-op MRI...');
-    % Create temporary folder
-    TmpDir = bst_get('BrainstormTmpDir', 0, 'resection_identification');
+    % Get temporary folder from container info
+    TmpDir = containerInfo.volumes{1,1};
     % Save pre-op MRI
     preOpNii = bst_fullfile(TmpDir, 'preop.nii');
     out_mri_nii(sMriPreOp, preOpNii);
     % Save post-op MRI
     postOpNii = bst_fullfile(TmpDir, 'postop.nii');
     sMriPostOp = in_mri_bst(MriFilePostOp);
     out_mri_nii(sMriPostOp, postOpNii);
-
-    % === CALL RESECTION-IDENTIFICATION PIPELINE ===
-    bst_progress('text', 'Calling resection-identification...');
-    % Get resection-identification executable
-    ResecExe = bst_fullfile(PlugDesc.Path, PlugDesc.SubFolder, PlugDesc.TestFile);
-    % Call resection-identification
-    strCall = ['"' ResecExe '"' ' ' '"' preOpNii '"' ' ' '"' postOpNii '"' ' ' '"' TmpDir '"'];
-    disp(['RESEC_ID > System call: ' strCall]);
-    tic;
-    status = system(strCall);
-    if (status ~= 0)
-        errMsg = 'Error during resection-identification, see logs in the command window.';
-        bst_progress('stop');
-        return;
+    % Make dir for processed files
+    mkdir(bst_fullfile(TmpDir, 'temp_dir_resection'));
+
+    % === RUN COMMAND IN CONTAINER =====
+    if isempty(errMsg) && containerInfo.isRunning
+        dataPath = containerInfo.volumes{1,2};
+        command = [' python3 auto_resection_mask.py ' dataPath '/preop.nii ' dataPath '/postop.nii'];
+        errMsg = bst_containers('ExecInContainer', containerInfo.name, command);
+    else
+        errMsg = 'Container is not running';
+    end
+    if ~isempty(errMsg)
+        return
     end
     disp(['RESEC_ID > Computation completed in: ' num2str(round(toc)) ' s']);
 
@@ -182,8 +185,10 @@
     Post2PreOpNii  = bst_fullfile(TmpDir, 'postop.nonlin.post2pre.nii.gz');
     MriFilePost2PreOp = import_mri(iSubject, Post2PreOpNii, 'Nifti1', 0, 1, 'postop_coreg_preop');
 
-    % Delete the temporary files
-    file_delete(TmpDir, 1, 1);
+    % Unload container plugin === Stop container and Delete bind files
+    if ensureRes > 0
+        bst_plugin('Unload', plugName);
+    end
     % Return success
     isOk = 1;
 end