make crops

Snakefile

 rule combine_tables:
 	input:
 		table1="{folder}/day1/table.csv",
-		table2="{folder}/day2/table.csv"
+		table2="{folder}/day2/table.csv",
+		crops = "{folder}/day2/BF_TRITC_aligned.crops.zarr"
 	params:
 		threshold="20" #minimum number of cells for day 2
 
@@ -26,6 +27,15 @@ rule align_and_count:
 	shell:
 		"python align.py {input.data} {output.zarr} {input.concentrations} {input.template} {input.labels} {output.table} 1" 
 
+rule make_crops:
+	input:
+		data = "{folder}/day{day}/BF_TRITC_aligned.zarr"
+	output:
+		zarr = directory("{folder}/day{day}/BF_TRITC_aligned.crops.zarr")
+	shell:
+		"python make_crops.py {input.data} {output.zarr}"
+
+
 rule correct_xy:
 	input:
 		zarr = "{folder}/day{day}/BF_TRITC_aligned-to-correct.zarr",

make_crops.py

+import pandas as pd
+import os
+import dask.array as da
+import re
+import fire
+
+EXPECTED_SHAPE = (6, 3, 7152, 22192)
+
+def main(aligned_path, out_path, size=300):
+    centers = pd.read_csv(
+        "v3/centers_bin16.csv", 
+        index_col=0
+    ).values * 8
+    _ = get_crops(aligned_path, out_path, centers=centers, size=size)
+    return 0
+
+def crop2d(img, center: tuple, size: int ):
+    im = img[
+        ...,
+        int(center[0]) - size // 2 : int(center[0]) + size // 2,
+        int(center[1]) - size // 2 : int(center[1]) + size // 2,
+    ]
+    return im
+
+def get_crops(path_to_zarr, out_path, centers, size):
+    if os.path.exists(out_path):
+        print(f"already exists, reading")
+        return da.from_zarr(out_path)
+    date = (re.compile("[0-9]{8}").findall(path_to_zarr)[0])
+    data = da.from_zarr(path_to_zarr+"/0/")
+    print(f"{date}: {data.shape} \n")
+    if not data.shape == EXPECTED_SHAPE:
+        return
+    crops = da.stack(map(lambda c: crop2d(data[:,:2], c, size), centers)).rechunk()
+    da.to_zarr(crops, out_path)
+    print(f"Saved {crops.shape} to {out_path}")
+    return crops
+
+
+if __name__ == "__main__":
+    fire.Fire(main)