|
5 | 5 | import logging |
6 | 6 | import os |
7 | 7 | import shutil |
| 8 | +import tempfile |
8 | 9 | from enum import Enum |
9 | 10 | from pathlib import Path |
10 | 11 | from typing import Dict, List |
11 | 12 |
|
12 | 13 | import h5py |
| 14 | +import nibabel as nib |
13 | 15 |
|
14 | 16 | from openlifu.nav.photoscan import Photoscan, load_data_from_photoscan |
15 | 17 | from openlifu.plan import Protocol, Run, Solution |
16 | 18 | from openlifu.util.json import PYFUSEncoder |
17 | 19 | from openlifu.util.types import PathLike |
| 20 | +from openlifu.util.volume_conversion import ( |
| 21 | + convert_dicom_to_nifti, |
| 22 | + is_dicom_file_or_directory, |
| 23 | +) |
18 | 24 | from openlifu.xdc import Transducer, TransducerArray |
19 | 25 | from openlifu.xdc.util import load_transducer_from_file |
20 | 26 |
|
|
24 | 30 |
|
25 | 31 | OnConflictOpts = Enum('OnConflictOpts', ['ERROR', 'OVERWRITE', 'SKIP']) |
26 | 32 |
|
| 33 | + |
27 | 34 | class Database: |
28 | 35 | def __init__(self, path: str | None = None): |
29 | 36 | if path is None: |
@@ -378,35 +385,57 @@ def write_volume(self, subject_id, volume_id, volume_name, volume_data_filepath, |
378 | 385 | if not Path(volume_data_filepath).exists(): |
379 | 386 | raise ValueError(f'Volume data filepath does not exist: {volume_data_filepath}') |
380 | 387 |
|
381 | | - volume_ids = self.get_volume_ids(subject_id) |
382 | | - if volume_id in volume_ids: |
383 | | - if on_conflict == OnConflictOpts.ERROR: |
384 | | - raise ValueError(f"Volume with ID {volume_id} already exists for subject {subject_id}.") |
385 | | - elif on_conflict == OnConflictOpts.OVERWRITE: |
386 | | - self.logger.info(f"Overwriting volume with ID {volume_id} for subject {subject_id}.") |
387 | | - elif on_conflict == OnConflictOpts.SKIP: |
388 | | - self.logger.info(f"Skipping volume with ID {volume_id} for subject {subject_id} as it already exists.") |
389 | | - return |
390 | | - else: |
391 | | - raise ValueError("Invalid 'on_conflict' option. Use 'error', 'overwrite', or 'skip'.") |
392 | | - |
393 | | - # Create volume metadata |
394 | | - volume_metadata_dict = {"id": volume_id, "name": volume_name, "data_filename": Path(volume_data_filepath).name} |
395 | | - volume_metadata_json = json.dumps(volume_metadata_dict, separators=(',', ':'), cls=PYFUSEncoder) |
396 | | - |
397 | | - # Save the volume metadata to a JSON file and copy volume data file to database |
398 | | - volume_metadata_filepath = self.get_volume_metadata_filepath(subject_id, volume_id) #subject_id/volume/volume_id/volume_id.json |
399 | | - Path(volume_metadata_filepath).parent.parent.mkdir(exist_ok=True) # volume directory |
400 | | - Path(volume_metadata_filepath).parent.mkdir(exist_ok=True) |
401 | | - with open(volume_metadata_filepath, 'w') as file: |
402 | | - file.write(volume_metadata_json) |
403 | | - shutil.copy(Path(volume_data_filepath), Path(volume_metadata_filepath).parent) |
404 | | - |
405 | | - if volume_id not in volume_ids: |
406 | | - volume_ids.append(volume_id) |
407 | | - self.write_volume_ids(subject_id, volume_ids) |
408 | | - |
409 | | - self.logger.info(f"Added volume with ID {volume_id} for subject {subject_id} to the database.") |
| 388 | + path = Path(volume_data_filepath) |
| 389 | + if path.is_dir() and not is_dicom_file_or_directory(volume_data_filepath): |
| 390 | + raise ValueError(f'Volume data filepath is a directory without DICOM files: {volume_data_filepath}') |
| 391 | + |
| 392 | + # convert dicom to nifti if needed |
| 393 | + temp_nifti_path = None |
| 394 | + if is_dicom_file_or_directory(volume_data_filepath): |
| 395 | + self.logger.info(f"Detected DICOM input for volume {volume_id}, converting to NIfTI format") |
| 396 | + with tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False) as temp_file: |
| 397 | + temp_nifti_path = Path(temp_file.name) |
| 398 | + |
| 399 | + try: |
| 400 | + convert_dicom_to_nifti(volume_data_filepath, temp_nifti_path) |
| 401 | + volume_data_filepath = temp_nifti_path |
| 402 | + except Exception as e: |
| 403 | + if temp_nifti_path.exists(): |
| 404 | + temp_nifti_path.unlink() |
| 405 | + raise RuntimeError(f"Failed to convert DICOM to NIfTI: {e}") from e |
| 406 | + |
| 407 | + try: |
| 408 | + volume_ids = self.get_volume_ids(subject_id) |
| 409 | + if volume_id in volume_ids: |
| 410 | + if on_conflict == OnConflictOpts.ERROR: |
| 411 | + raise ValueError(f"Volume with ID {volume_id} already exists for subject {subject_id}.") |
| 412 | + elif on_conflict == OnConflictOpts.OVERWRITE: |
| 413 | + self.logger.info(f"Overwriting volume with ID {volume_id} for subject {subject_id}.") |
| 414 | + elif on_conflict == OnConflictOpts.SKIP: |
| 415 | + self.logger.info(f"Skipping volume with ID {volume_id} for subject {subject_id} as it already exists.") |
| 416 | + return |
| 417 | + else: |
| 418 | + raise ValueError("Invalid 'on_conflict' option. Use 'error', 'overwrite', or 'skip'.") |
| 419 | + |
| 420 | + volume_metadata_dict = {"id": volume_id, "name": volume_name, "data_filename": Path(volume_data_filepath).name} |
| 421 | + volume_metadata_json = json.dumps(volume_metadata_dict, separators=(',', ':'), cls=PYFUSEncoder) |
| 422 | + |
| 423 | + volume_metadata_filepath = self.get_volume_metadata_filepath(subject_id, volume_id) |
| 424 | + Path(volume_metadata_filepath).parent.parent.mkdir(exist_ok=True) |
| 425 | + Path(volume_metadata_filepath).parent.mkdir(exist_ok=True) |
| 426 | + with open(volume_metadata_filepath, 'w') as file: |
| 427 | + file.write(volume_metadata_json) |
| 428 | + shutil.copy(Path(volume_data_filepath), Path(volume_metadata_filepath).parent) |
| 429 | + |
| 430 | + if volume_id not in volume_ids: |
| 431 | + volume_ids.append(volume_id) |
| 432 | + self.write_volume_ids(subject_id, volume_ids) |
| 433 | + |
| 434 | + self.logger.info(f"Added volume with ID {volume_id} for subject {subject_id} to the database.") |
| 435 | + finally: |
| 436 | + # cleanup temp nifti file |
| 437 | + if temp_nifti_path is not None and temp_nifti_path.exists(): |
| 438 | + temp_nifti_path.unlink() |
410 | 439 |
|
411 | 440 | def write_photocollection(self, subject_id, session_id, reference_number: str, photo_paths: List[PathLike], on_conflict=OnConflictOpts.ERROR): |
412 | 441 | """ Writes a photocollection to database and copies the associated |
@@ -980,7 +1009,6 @@ def load_volume(self, subject, volume_id): |
980 | 1009 | if not volume_data_filepath.exists() or not volume_data_filepath.is_file(): |
981 | 1010 | self.logger.error(f"Volume data file not found for volume {volume_id}, subject {subject.id}") |
982 | 1011 | raise FileNotFoundError(f"Volume data file not found for volume {volume_id}, subject {subject.id}") |
983 | | - import nibabel as nib |
984 | 1012 | # Load the volume data using nibabel |
985 | 1013 | volume_data = nib.load(volume_data_filepath) |
986 | 1014 | self.logger.info(f"Loaded volume {volume_id} for subject {subject.id}") |
|
0 commit comments