Update doc string

src/lazyslide/_preprocess/_graph.py

@@ -25,6 +25,28 @@ def tile_graph(
     tile_key: str = Key.tiles,
     table_key: str = None,
 ):
+    """
+    Compute the spatial graph of the tiles.
+
+    Parameters
+    ----------
+    wsi : :class:`WSIData <wsidata.WSIData>`
+        The WSIData object to work on.
+    n_neighs : int, default: 6
+        The number of neighbors to consider.
+    n_rings : int, default: 1
+        The number of rings to consider.
+    delaunay : bool, default: False
+        Whether to use Delaunay triangulation.
+    transform : str, default: None
+        The transformation to apply to the graph.
+    set_diag : bool, default: False
+        Whether to set the diagonal to 1.
+    tile_key : str, default: 'tiles'
+        The tile key.
+    table_key : str, default: None
+        The table key to store the graph.
+    """
     coords = wsi.sdata[tile_key][["x", "y"]].values
     Adj, Dst = _spatial_neighbor(
         coords, n_neighs, delaunay, n_rings, transform, set_diag

src/lazyslide/_preprocess/_load.py

@@ -20,7 +20,28 @@ def load_annotations(
     min_area: float = 1e2,
     key_added: str = "annotations",
 ):
-    """Load the geojson file and add it to the WSI data"""
+    """Load the annotation file and add it to the WSIData
+
+    Parameters
+    ----------
+    wsi : :class:`WSIData <wsidata.WSIData>`
+        The WSIData object to work on.
+    annotations : str, Path, GeoDataFrame
+        The path to the annotation file or the GeoDataFrame.
+    explode : bool, default: True
+        Whether to explode the annotations.
+    in_bounds : bool, default: False
+        Whether to move the annotations to the slide bounds.
+    join_with : str, List[str], default: 'tissues'
+        The key to join the annotations with.
+    join_to : str, default: None
+        The key to join the annotations to.
+    min_area : float, default: 1e2
+        The minimum area of the annotation.
+    key_added : str, default: 'annotations'
+        The key to store the annotations.
+
+    """
     import geopandas as gpd
 
     if isinstance(annotations, (str, Path)):

src/lazyslide/_preprocess/_tiles.py

@@ -36,7 +36,8 @@ def tile_tissues(
 
     Parameters
     ----------
-    wsi : {wsi}
+    wsi : :class:`WSIData <wsidata.WSIData>`
+        The WSIData object to work on.
     tile_px : int, (int, int)
         The size of the tile, if tuple, (W, H).
     stride_px : int, (int, int), default: None
@@ -229,10 +230,14 @@ def tile_tissues(
         level=ops_level,
         downsample=downsample,
         mpp=mpp,
-        height=tile_h,
-        width=tile_w,
-        raw_height=ops_tile_h,
-        raw_width=ops_tile_w,
+        height=int(tile_h),
+        width=int(tile_w),
+        stride_height=int(stride_h),
+        stride_width=int(stride_w),
+        raw_height=int(ops_tile_h),
+        raw_width=int(ops_tile_w),
+        raw_stride_height=int(ops_stride_h),
+        raw_stride_width=int(ops_stride_w),
         tissue_name=tissue_key,
     )
 
@@ -266,8 +271,8 @@ def tiles_qc(
 
     Parameters
     ----------
-    wsi : WSI
-        The whole slide image object.
+    wsi : :class:`WSIData <wsidata.WSIData>`
+        The WSIData object to work on.
     scorers : Scorer
         The scorer object or a callable that takes in an image and returns a score.
         You can also pass in a string:

src/lazyslide/_tools/_features.py

@@ -266,10 +266,10 @@ def feature_extraction(
             features = []
             with torch.inference_mode():
                 for batch in loader:
-                    batch = batch.to(device)
-                    output = model(batch)
+                    image = batch["image"].to(device)
+                    output = model(image)
                     features.append(output.cpu().numpy())
-                    progress_bar.update(task, advance=len(batch))
+                    progress_bar.update(task, advance=len(image))
                     del batch  # Free up memory
             # The progress bar may not reach 100% if exit too early
             # Force update
@@ -287,7 +287,7 @@ def _inference(dataset, chunk, model, queue):
     with torch.inference_mode():
         X = []
         for c in chunk:
-            img = dataset[c]
+            img = dataset[c]["image"]
             # image to 4d
             img = img.unsqueeze(0)
             output = model(img)