Add feature "include dicom files"

plugins/slicer/MONAILabel/MONAILabel.py

@@ -20,6 +20,7 @@
 from urllib.parse import quote_plus
 
 import ctk
+import DICOMScalarVolumePlugin
 import qt
 import SampleData
 import SimpleITK as sitk
@@ -190,6 +191,19 @@ def __init__(self, parent):
             str(qt.SIGNAL("valueAsIntChanged(int)")),
         )
 
+        includeDicomFilesCheckBox = qt.QCheckBox()
+        includeDicomFilesCheckBox.checked = False
+        includeDicomFilesCheckBox.toolTip = _(
+            "Enable this option to include dicom files in server-client data exchange"
+        )
+        groupLayout.addRow(_("Include DICOM files:"), includeDicomFilesCheckBox)
+        parent.registerProperty(
+            "MONAILabel/includeDicomFiles",
+            ctk.ctkBooleanMapper(includeDicomFilesCheckBox, "checked", str(qt.SIGNAL("toggled(bool)"))),
+            "valueAsInt",
+            str(qt.SIGNAL("valueAsIntChanged(int)")),
+        )
+
         vBoxLayout.addWidget(groupBox)
         vBoxLayout.addStretch(1)
 
@@ -1369,6 +1383,14 @@ def getPermissionForImageDataUpload(self):
             dontShowAgainSettingsKey="MONAILabel/showImageDataSendWarning",
         )
 
+    def get_dicom_files(self):
+        dicom_files = []
+        for path, subdir, files in os.walk(self.tmpdir):
+            for file in files:
+                if file[-3:] == "dcm":
+                    dicom_files.append(os.path.join(path, file))
+        return dicom_files
+
     def onUploadImage(self, init_sample=True, session=False):
         volumeNode = slicer.mrmlScene.GetFirstNodeByClass("vtkMRMLScalarVolumeNode")
         image_id = volumeNode.GetName()
@@ -1384,12 +1406,40 @@ def onUploadImage(self, init_sample=True, session=False):
             start = time.time()
             slicer.util.saveNode(volumeNode, in_file)
             logging.info(f"Saved Input Node into {in_file} in {time.time() - start:3.1f}s")
-            self.reportProgress(30)
+            last_report_progress = 30
+            self.reportProgress(last_report_progress)
+
+            # if includeDicomFilesCheckBox is marked, save original dicom files on the server
+            if slicer.util.settingsValue("MONAILabel/includeDicomFiles", False, converter=slicer.util.toBool):
+                start = time.time()
+                shNode = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
+                volumeShItemID = shNode.GetItemByDataNode(volumeNode)
+
+                exporter = DICOMScalarVolumePlugin.DICOMScalarVolumePluginClass()
+                exportables = exporter.examineForExport(volumeShItemID)
+                for exp in exportables:
+                    exp.directory = self.tmpdir
+                exporter.export(exportables)
+                logging.info(f"Saved dicom files in {time.time() - start:3.1f}s")
+                last_report_progress = 50
+                self.reportProgress(last_report_progress)
 
             if session:
                 self.current_sample["session_id"] = self.logic.create_session(in_file)["session_id"]
             else:
                 self.logic.upload_image(in_file, image_id)
+
+                # if includeDicomFilesCheckBox is marked, save original dicom files on the client
+                if slicer.util.settingsValue("MONAILabel/includeDicomFiles", False, converter=slicer.util.toBool):
+                    dcm_filenames = self.get_dicom_files()
+                    report_progress_increment = round(last_report_progress / len(dcm_filenames), 1)
+                    for dcm_fullpath in dcm_filenames:
+                        dcm_filename = dcm_fullpath.split("/")[-1]
+                        dcm_filename = ".".join(dcm_filename.split(".")[:-1])
+                        self.logic.upload_image(dcm_fullpath, dcm_filename)
+                        last_report_progress += report_progress_increment
+                        self.reportProgress(last_report_progress)
+
                 self.current_sample["session"] = False
             self.reportProgress(100)
 
@@ -1485,7 +1535,11 @@ def onSaveLabel(self):
                 self.reportProgress(30)
 
             self.updateServerSettings()
-            result = self.logic.save_label(self.current_sample["id"], label_in, {"label_info": label_info})
+            result = self.logic.save_label(
+                self.current_sample["id"],
+                label_in,
+                {"label_info": label_info, "model": model},
+            )
             self.fetchInfo()
 
             if slicer.util.settingsValue("MONAILabel/autoUpdateModelV2", False, converter=slicer.util.toBool):
@@ -1531,6 +1585,10 @@ def onClickSegmentation(self):
         if not self.current_sample:
             return
 
+        if self.current_sample.get("session"):
+            if not self.onUploadImage(init_sample=False):
+                return
+
         start = time.time()
         result_file = None
         try: