Skip to content

Commit

Permalink
handle unsupported file formats, improve error handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@almazgimaev
almazgimaev committed Nov 13, 2023
1 parent 1c8de1d commit f0f8429
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 74 deletions.
144 changes: 72 additions & 72 deletions src/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,86 +10,86 @@ def import_dicom_volumes(
api: sly.Api, task_id: int, context: dict, state: dict, app_logger
) -> None:
save_path = f.download_data_from_team_files(api=api, task_id=task_id, save_path=g.STORAGE_DIR)
project_dir = f.get_project_dir(path=save_path)

if g.PROJECT_ID is None:
project_name = (
f.get_project_name_from_input_path(project_dir)
if len(g.OUTPUT_PROJECT_NAME) == 0
else g.OUTPUT_PROJECT_NAME
)

sly.logger.debug(f"Project name: {project_name}")

project = g.api.project.create(
workspace_id=g.WORKSPACE_ID,
name=project_name,
type=sly.ProjectType.VOLUMES,
change_name_if_conflict=True,
)
else:
project = api.project.get_info_by_id(g.PROJECT_ID)

if g.DATASET_ID is None:
dataset = g.api.dataset.create(
project_id=project.id, name=g.DEFAULT_DATASET_NAME, change_name_if_conflict=True
)
used_volumes_names = []
else:
dataset = api.dataset.get_info_by_id(g.DATASET_ID)
used_volumes_names = [volume.name for volume in api.volume.get_list(dataset.id)]

# DICOM
series_infos = sly.volume.inspect_dicom_series(root_dir=project_dir)
# NRRD
nrrd_paths = sly.volume.inspect_nrrd_series(root_dir=project_dir)

if len(series_infos) == 0 and len(nrrd_paths) == 0:
sly.logger.warn("No volumes were found. Please, check your input directory.")
else:
if save_path:
project_dir = f.get_project_dir(path=save_path)

if g.PROJECT_ID is None:
project_name = (
f.get_project_name_from_input_path(project_dir)
if len(g.OUTPUT_PROJECT_NAME) == 0
else g.OUTPUT_PROJECT_NAME
)

# DICOM
for serie_id, files in series_infos.items():
item_path = files[0]
if sly.volume.get_extension(path=item_path) is None:
sly.logger.warn(
f"Can not recognize file extension {item_path}, serie will be skipped"
)
continue
name = f"{serie_id}.nrrd"
name = f.generate_free_name(
used_names=used_volumes_names, possible_name=name, with_ext=True
sly.logger.debug(f"Project name: {project_name}")

project = g.api.project.create(
workspace_id=g.WORKSPACE_ID,
name=project_name,
type=sly.ProjectType.VOLUMES,
change_name_if_conflict=True,
)
used_volumes_names.append(name)
g.api.volume.upload_dicom_serie_paths(
dataset_id=dataset.id,
name=name,
paths=files,
log_progress=True,
anonymize=g.ANONYMIZE_VOLUMES,
else:
project = api.project.get_info_by_id(g.PROJECT_ID)

if g.DATASET_ID is None:
dataset = g.api.dataset.create(
project_id=project.id, name=g.DEFAULT_DATASET_NAME, change_name_if_conflict=True
)
used_volumes_names = []
else:
dataset = api.dataset.get_info_by_id(g.DATASET_ID)
used_volumes_names = [volume.name for volume in api.volume.get_list(dataset.id)]

# DICOM
series_infos = sly.volume.inspect_dicom_series(root_dir=project_dir)
# NRRD
for nrrd_path in nrrd_paths:
name = sly.fs.get_file_name_with_ext(path=nrrd_path)
name = f.generate_free_name(
used_names=used_volumes_names, possible_name=name, with_ext=True
)
used_volumes_names.append(name)
g.api.volume.upload_nrrd_serie_path(
dataset_id=dataset.id, name=name, path=nrrd_path, log_progress=True
)
nrrd_paths = sly.volume.inspect_nrrd_series(root_dir=project_dir)

if g.REMOVE_SOURCE and not g.IS_ON_AGENT:
if g.INPUT_DIR is not None:
path_to_remove = g.INPUT_DIR
if len(series_infos) == 0 and len(nrrd_paths) == 0:
sly.logger.warn("No volumes were found. Please, check your input directory.")
else:
path_to_remove = g.INPUT_FILE
api.file.remove(team_id=g.TEAM_ID, path=path_to_remove)
source_dir_name = path_to_remove.lstrip("/").rstrip("/")
sly.logger.info(msg=f"Source directory: '{source_dir_name}' was successfully removed.")
# DICOM
for serie_id, files in series_infos.items():
item_path = files[0]
if sly.volume.get_extension(path=item_path) is None:
sly.logger.warn(
f"Can not recognize file extension {item_path}, serie will be skipped"
)
continue
name = f"{serie_id}.nrrd"
name = f.generate_free_name(
used_names=used_volumes_names, possible_name=name, with_ext=True
)
used_volumes_names.append(name)
g.api.volume.upload_dicom_serie_paths(
dataset_id=dataset.id,
name=name,
paths=files,
log_progress=True,
anonymize=g.ANONYMIZE_VOLUMES,
)

# NRRD
for nrrd_path in nrrd_paths:
name = sly.fs.get_file_name_with_ext(path=nrrd_path)
name = f.generate_free_name(
used_names=used_volumes_names, possible_name=name, with_ext=True
)
used_volumes_names.append(name)
g.api.volume.upload_nrrd_serie_path(
dataset_id=dataset.id, name=name, path=nrrd_path, log_progress=True
)

if g.REMOVE_SOURCE and not g.IS_ON_AGENT:
if g.INPUT_DIR is not None:
path_to_remove = g.INPUT_DIR
else:
path_to_remove = g.INPUT_FILE
api.file.remove(team_id=g.TEAM_ID, path=path_to_remove)
source_dir_name = path_to_remove.lstrip("/").rstrip("/")
sly.logger.info(msg=f"Source directory: '{source_dir_name}' was successfully removed.")

api.task.set_output_project(task_id, project.id, project.name)
api.task.set_output_project(task_id, project.id, project.name)
g.my_app.stop()


Expand Down
19 changes: 17 additions & 2 deletions src/sly_functions.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import functools
import os
import shutil
from typing import Callable
import tarfile
import zipfile
Expand Down Expand Up @@ -119,6 +118,17 @@ def download_data_from_team_files(api: sly.Api, task_id: int, save_path: str) ->
save_archive_path = local_save_path
sly.fs.unpack_archive(save_archive_path, extrack_dir)
silent_remove(save_archive_path)
elif local_save_path.endswith(".nii.gz") or local_save_path.endswith(".nii"):
file_name = get_file_name_with_ext(local_save_path)
msg = f"Nifti files are not supported. Please, upload .nrrd or .dcm files archive or directory."
sly.logger.error(msg, exc_info=False)
api.task.set_output_error(task_id, msg, description=f"File: {file_name}")
return None
elif local_save_path.endswith(".gz") and not is_archive(local_save_path):
msg = "Unsupported archive format. Please, use .zip, .tar or .tar.gz files."
sly.logger.error(msg, exc_info=False)
api.task.set_output_error(task_id, msg, description=f"File: {local_save_path}")
return None

return save_path

Expand All @@ -140,14 +150,19 @@ def generate_free_name(used_names, possible_name, with_ext=False, extend_used_na
used_names.add(res_name)
return res_name


def get_project_dir(path: str) -> str:
"""Returns project directory."""

def _volumes_exists(path: str) -> bool:
"""Returns True if path contains volumes."""
listdir = sly.fs.list_files(path)
if len([f for f in listdir if sly.volume.get_extension(path=f) is not None]) > 0:
return True
return False

all_volume_dirs = [d for d in sly.fs.dirs_filter(path, _volumes_exists)]
if len(all_volume_dirs) == 0:
return path
common_prefix = os.path.commonprefix(all_volume_dirs)
return common_prefix
return common_prefix

0 comments on commit f0f8429

Please sign in to comment.