Skip to content

Commit 1467969

Browse files
fix(libcommon): address CodeRabbit nitpick for Nifti handler
- Use hasattr(value, 'to_bytes') instead of isinstance check to support Nifti1Image, Nifti2Image, and other nibabel types - Preserve original file extension (.nii vs .nii.gz) - Remove unused nibabel import
1 parent b244a6d commit 1467969

1 file changed

Lines changed: 13 additions & 6 deletions

File tree

libs/libcommon/src/libcommon/viewer_utils/features.py

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -364,11 +364,18 @@ def nifti(
364364
storage_client: StorageClient,
365365
json_path: Optional[list[Union[str, int]]] = None,
366366
) -> Any:
367-
import nibabel as nib
368-
369367
if value is None:
370368
return None
371369

370+
# Determine file extension (preserve .nii vs .nii.gz)
371+
file_ext = ".nii.gz" # default
372+
if isinstance(value, dict) and "path" in value and isinstance(value["path"], str):
373+
path = value["path"]
374+
if path.endswith(".nii.gz"):
375+
file_ext = ".nii.gz"
376+
elif path.endswith(".nii"):
377+
file_ext = ".nii"
378+
372379
# Extract bytes from various input formats
373380
if isinstance(value, dict):
374381
if value.get("bytes"):
@@ -380,12 +387,12 @@ def nifti(
380387
raise ValueError(f"Nifti value must have 'bytes' or valid 'path': {value}")
381388
elif isinstance(value, bytes):
382389
nifti_bytes = value
383-
elif isinstance(value, nib.nifti1.Nifti1Image):
384-
# Handle decoded nibabel image - convert back to bytes
390+
elif hasattr(value, "to_bytes"):
391+
# Handle any nibabel image type (Nifti1Image, Nifti2Image, etc.)
385392
nifti_bytes = value.to_bytes()
386393
else:
387394
raise TypeError(
388-
"Nifti cell must be bytes, an encoded dict of a Nifti file, or a nibabel Nifti1Image, "
395+
"Nifti cell must be bytes, an encoded dict of a Nifti file, or a nibabel image, "
389396
f"but got {str(value)[:300]}{'...' if len(str(value)) > 300 else ''}"
390397
)
391398

@@ -396,7 +403,7 @@ def nifti(
396403
split=split,
397404
row_idx=row_idx,
398405
column=featureName,
399-
filename=f"{append_hash_suffix('nifti', json_path)}.nii.gz",
406+
filename=f"{append_hash_suffix('nifti', json_path)}{file_ext}",
400407
nifti_file_bytes=nifti_bytes,
401408
storage_client=storage_client,
402409
)

0 commit comments

Comments
 (0)