Feature key renaming and set a maximum size for datasets (#103)

* Add feature key renaming

* Set a maximum size for datasets

Author: LouisK92

src/data_processors/process_dataset/script.py

@@ -13,6 +13,50 @@
 }
 ### VIASH END
 
+
+def get_crop_coords(sdata, max_n_pixels=50000*50000):
+    """Get the crop coordinates to subset the sdata to max_n_pixels
+    
+    Arguments
+    ---------
+    sdata: spatialdata.SpatialData
+        The spatial data to crop
+    max_n_pixels: int
+        The maximum number of pixels to keep
+        
+    Returns
+    -------
+    crop_coords: dict
+        The crop coordinates
+    """
+    
+    _, h, w = sdata['morphology_mip']["scale0"].image.shape
+    #h, w = sdata
+    
+    # Check if the image is already below the maximum number of pixels
+    if h * w <= max_n_pixels:
+        return None
+    
+    # Initialize with square crop
+    h_crop = w_crop = int(np.sqrt(max_n_pixels))
+    
+    # Adjust crop if necessary to fit the image
+    if h_crop > h:
+        h_crop = h
+        w_crop = int(max_n_pixels / h_crop)
+    elif w_crop > w:
+        w_crop = w
+        h_crop = int(max_n_pixels / w_crop)
+        
+    # Center the crop
+    h_offset = (h - h_crop) // 2
+    w_offset = (w - w_crop) // 2
+        
+    crop = [[h_offset, h_offset + h_crop], [w_offset, w_offset + w_crop]]
+        
+    return crop
+
+
 # Load the single-cell data
 adata = ad.read_h5ad(par["input_sc"])
 
@@ -41,6 +85,26 @@
         sdata.table.uns[orig_col] = sdata.table.uns[col]
     sdata.table.uns[col] = par[col]
 
+# Correct the feature_key attribute in sdata if needed
+# NOTE: it would have been better to do this in the loader scripts, but this way the datasets don't need to be re-downloaded
+if "feature_key" in sdata['transcripts'].attrs["spatialdata_attrs"]:
+    feature_key = sdata['transcripts'].attrs["spatialdata_attrs"]["feature_key"]
+    if feature_key != "feature_name":
+        sdata['transcripts'].attrs["spatialdata_attrs"]["feature_key"] = "feature_name"
+
+# Crop datasets that are too large
+crop_coords = get_crop_coords(sdata)
+if crop_coords is not None:
+    sdata_output = sdata.query.bounding_box(
+        axes=["y", "x"],
+        min_coordinate=[crop_coords[0][0], crop_coords[1][0]],
+        max_coordinate=[crop_coords[0][1], crop_coords[1][1]],
+        target_coordinate_system="global",
+        filter_table=True,
+    )
+else:
+    sdata_output = sdata
+    
 # Save the single-cell data
 adata.write_h5ad(par["output_sc"], compression="gzip")
 
@@ -49,4 +113,4 @@
     shutil.rmtree(par["output_sp"])
 
 # Save the spatial data
-sdata.write(par["output_sp"], overwrite=True)
+sdata_output.write(par["output_sp"], overwrite=True)