From f86d87b795408a667d15ebd4c0f68d8a790dcefc Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Thu, 13 Aug 2026 04:00:42 -0700 Subject: [PATCH] user_dataset: gzip uncompressed uploads on store _promote used shutil.copy2 into a .nii.gz name, so a raw .nii upload was stored uncompressed under a .nii.gz name -- wasting space and unreadable (nibabel gunzips by extension). Add _store_nifti_gz: preserve bytes when the source is already gzip (the usual .nii.gz), otherwise gzip it. Applied to both the CT and combined_labels. Test: a raw .nii upload is stored as genuine gzip and loads back. 12/12 pass. --- flask-server/services/user_dataset.py | 20 ++++++++++++++-- flask-server/tests/unit/test_user_dataset.py | 25 ++++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/flask-server/services/user_dataset.py b/flask-server/services/user_dataset.py index d05c7da0..b88f65ca 100644 --- a/flask-server/services/user_dataset.py +++ b/flask-server/services/user_dataset.py @@ -313,6 +313,22 @@ def _write_sublabels(combined_labels_path: str, seg_dir: str, existing_seg_dir: return written +def _store_nifti_gz(src: str, dst: str) -> None: + """Store a NIfTI at dst as a genuinely gzip-compressed .nii.gz. If src is + already gzip (the usual .nii.gz upload) its bytes are preserved exactly; a raw + .nii is gzipped on the way in. Without this, copying an uncompressed upload to + a .nii.gz name would waste space AND write a file nibabel can't load (it + gunzips by extension).""" + with open(src, "rb") as f: + is_gzip = f.read(2) == b"\x1f\x8b" + if is_gzip: + shutil.copy2(src, dst) + else: + import gzip + with open(src, "rb") as fin, gzip.open(dst, "wb") as fout: + shutil.copyfileobj(fin, fout) + + def _promote(root: str, case_id: str, ct_path: str, combined_labels_path: str, existing_seg_dir: Optional[str], metadata: dict) -> None: """Write all of a case's files into `.partial` staging dirs, then publish each @@ -327,8 +343,8 @@ def _promote(root: str, case_id: str, ct_path: str, combined_labels_path: str, try: os.makedirs(img_stage, exist_ok=True) os.makedirs(mask_stage, exist_ok=True) - shutil.copy2(ct_path, os.path.join(img_stage, "ct.nii.gz")) - shutil.copy2(combined_labels_path, os.path.join(mask_stage, "combined_labels.nii.gz")) + _store_nifti_gz(ct_path, os.path.join(img_stage, "ct.nii.gz")) + _store_nifti_gz(combined_labels_path, os.path.join(mask_stage, "combined_labels.nii.gz")) organs = _write_sublabels(combined_labels_path, os.path.join(mask_stage, "segmentations"), existing_seg_dir) meta = {**metadata, "case_id": case_id, "organs": organs} diff --git a/flask-server/tests/unit/test_user_dataset.py b/flask-server/tests/unit/test_user_dataset.py index b1021979..a745c554 100644 --- a/flask-server/tests/unit/test_user_dataset.py +++ b/flask-server/tests/unit/test_user_dataset.py @@ -211,6 +211,31 @@ def test_voxel_guard_blocks_oom(ud, tmp_path, monkeypatch): assert not ok and reason.startswith("too_many_voxels") +def test_uncompressed_upload_stored_as_gzip(ud, tmp_path): + """A raw .nii upload must be gzipped on store, not copied under a .nii.gz name + (which would be uncompressed and unloadable).""" + np = pytest.importorskip("numpy") + nib = pytest.importorskip("nibabel") + vol = np.full((64, 64, 64), -1000.0, dtype="float32") + vol[20:40, 20:40, 20:40] = 60.0 + vol[30:34, 30:34, 30:34] = 400.0 + ct = str(tmp_path / "scan.nii") # RAW, uncompressed upload + nib.save(nib.Nifti1Image(vol, np.eye(4)), ct) + with open(ct, "rb") as f: + assert f.read(2) != b"\x1f\x8b" # confirm the source really is raw + + out = tmp_path / "out"; out.mkdir() + mask = np.zeros((64, 64, 64), "uint8"); mask[10:30, 10:30, 10:30] = 14; mask[35:45, 35:45, 35:45] = 17 + _write_nifti(str(out / "combined_labels.nii.gz"), mask) + + ud._admit_and_store(ct, str(out), "ePAI", "u1", "1.2.3.4", "s") + stored = os.path.join(ud._root(), "image_only", "USER_00000001", "ct.nii.gz") + assert os.path.exists(stored) + with open(stored, "rb") as f: + assert f.read(2) == b"\x1f\x8b" # stored file is genuinely gzip + assert np.asarray(nib.load(stored).dataobj).shape == (64, 64, 64) # and loads back + + def test_rejects_non_finite(ud, tmp_path): np = pytest.importorskip("numpy") pytest.importorskip("nibabel")