cellmap-challenge-analysis/fix_dtype.py at main · janelia-cellmap/cellmap-challenge-analysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import zarr
import json
from upath import UPath

from tqdm.auto import tqdm
from cellmap_segmentation_challenge.utils.eval_utils.submission import (
    _prepare_submission,
)

dtype_mappings = {
    "bool": "uint8",
}


def fix_dtype(zarr_path):
    _ = _prepare_submission(zarr_path.replace(".zarr", ""))
    zarr_group = zarr.open(zarr_path, mode="r+")
    for crop, group in tqdm(
        zarr_group.items(),
        desc=f"Processing {zarr_path} crops",
        unit="crop",
        dynamic_ncols=True,
        leave=True,
        position=0,
    ):
        p_bar = tqdm(
            group.items(),
            dynamic_ncols=True,
            leave=False,
            desc=f"Processing {crop} organelles",
            unit="organelle",
            position=1,
        )
        for organelle, array in p_bar:
            original_dtype = str(array.dtype)
            # print(f"Checking dtype for {crop}/{organelle}: {original_dtype}")
            if original_dtype in dtype_mappings:
                new_dtype = dtype_mappings[original_dtype]
                p_bar.set_postfix_str(
                    f"Fixing dtype for {crop}/{organelle}: {original_dtype} -> {new_dtype}"
                )
                data = array[:]
                array[:] = data.astype(new_dtype)

                # Load existing metadata
                metadata = json.load(
                    open(UPath(zarr_path) / crop / organelle / ".zarray")
                )

                # Update dtype and fill_value in metadata
                metadata["dtype"] = "|u1"
                metadata["fill_value"] = 0

                # Save updated metadata back to .zarray
                with open(UPath(zarr_path) / crop / organelle / ".zarray", "w") as f:
                    json.dump(metadata, f, indent=4)

    print("Successfully fixed dtypes where necessary.")


def fix_metadata(zarr_path):
    zarr_group = zarr.open(zarr_path, mode="r+")
    for crop, group in tqdm(
        zarr_group.items(),
        desc=f"Processing {zarr_path} crops",
        unit="crop",
        dynamic_ncols=True,
        leave=True,
        position=0,
    ):
        p_bar = tqdm(
            group.keys(),
            dynamic_ncols=True,
            leave=False,
            desc=f"Processing {crop} organelles",
            unit="organelle",
            position=1,
        )
        for organelle in p_bar:
            # Load existing metadata
            metadata = json.load(open(UPath(zarr_path) / crop / organelle / ".zarray"))

            # Update dtype and fill_value in metadata
            metadata["dtype"] = "|u1"
            metadata["fill_value"] = 0

            # Save updated metadata back to .zarray
            with open(UPath(zarr_path) / crop / organelle / ".zarray", "w") as f:
                json.dump(metadata, f, indent=4)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Fix dtype in Zarr files.")
    parser.add_argument(
        "paths",
        nargs="*",
        default=[],
        help="Paths to Zarr files to fix (can be multiple)",
    )
    parser.add_argument(
        "--metadata-only",
        action="store_true",
        help="Only fix metadata without changing the actual data type",
    )
    parser.add_argument(
        "--csv_path",
        type=str,
        help="Path to a CSV file containing Zarr paths",
        default=None,
    )

    args = parser.parse_args()
    paths = args.paths
    if args.csv_path:
        import pandas as pd

        df = pd.read_csv(args.csv_path)
        csv_paths = df["data_path"].tolist()
        csv_paths = [UPath(p).with_suffix(".zarr").path for p in csv_paths]
        paths.extend(csv_paths)

    for zarr_path in paths:
        try:
            if args.metadata_only:
                fix_metadata(zarr_path)
            else:
                fix_dtype(zarr_path)
        except Exception as e:
            print(f"An error occurred: {e}")