BioPredict/1_patch_encoder.py at main · AllenSpike/BioPredict · 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
import os
import json
import argparse

os.environ['SF_SLIDE_BACKEND'] = 'libvips'
os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = str(pow(2, 50))

import numpy as np
import slideflow as sf
from slideflow.slide import qc


def parse_args():
    parser = argparse.ArgumentParser(
        description="Step 1: patching tiling and feature extraction from WSIs with slideflow."
    )
    parser.add_argument('--project_root', required=True,
                        help='Slideflow project directory (created if not existing).')
    parser.add_argument('--slides', required=True,
                        help='Directory containing whole-slide images.')
    parser.add_argument('--annotation', required=True,
                        help='Directory for sample annotation CSV.')
    parser.add_argument('--roi_dest', default=None,
                        help='Directory for ROI CSV files (optional).')
    parser.add_argument('--tile_px', type=int, default=224,
                        help='Tile size in pixels (default: 224).')
    parser.add_argument('--mpp', type=float, default=0.5,
                        help='Microns per pixel of the slides; tile_um = round(tile_px * mpp). '
                             'Ignored if --tile_um is set. (default: 0.5)')
    parser.add_argument('--tile_um', type=int, default=None,
                        help='Tile size in microns. Overrides --mpp if provided.')
    parser.add_argument('--encoder', default='uni',
                        help="Feature encoder name, e.g. 'uni', 'ctranspath', 'plip', 'phikon' (default: uni).")
    parser.add_argument('--weights', required=True,
                        help='Path to local feature-encoder weights.')
    parser.add_argument('--output_dir', required=True,
                        help='Output directory for feature bags (.pt files).')
    parser.add_argument('--normalizer', default='reinhard_fast',
                        help='Stain normalization method (default: reinhard_fast).')
    parser.add_argument('--num_threads', type=int, default=10,
                        help='Number of extraction threads (default: 10).')
    parser.add_argument('--geojson_dir', default=None,
                        help='Directory for per-slide patch-coordinate GeoJSON files '
                             '(default: <output_dir>/geojson).')
    return parser.parse_args()


def save_patch_geojson(outdir, geojson_dir, tile_px):
    """Save patch coordinates of each WSI as a GeoJSON of square polygons."""
    os.makedirs(geojson_dir, exist_ok=True)

    for fname in sorted(os.listdir(outdir)):
        if not fname.endswith('.index.npz'):
            continue

        slide = fname[:-len('.index.npz')]
        coords = np.load(os.path.join(outdir, fname))['arr_0']
        if coords.shape[0] == 0:
            continue

        # estimate patch side length from coordinate spacing
        uniq_x = np.unique(coords[:, 0])
        gaps = np.diff(uniq_x)
        side = float(gaps[gaps > 0].min()) if np.any(gaps > 0) else float(tile_px)

        features = []
        for x, y in coords:
            x, y = float(x), float(y)
            ring = [[x, y], [x + side, y], [x + side, y + side], [x, y + side], [x, y]]
            features.append({
                "type": "Feature",
                "properties": {"slide": slide, "x": x, "y": y},
                "geometry": {"type": "Polygon", "coordinates": [ring]},
            })

        fc = {"type": "FeatureCollection", "features": features}
        with open(os.path.join(geojson_dir, slide + '.geojson'), 'w') as f:
            json.dump(fc, f)


def main():
    args = parse_args()

    tile_um = args.tile_um if args.tile_um is not None else round(args.tile_px * args.mpp)

    if not os.path.exists(os.path.join(args.project_root, 'settings.json')):
        project = sf.create_project(
            root=args.project_root,
            slides=args.slides,
            roi_dest=args.roi_dest,
            annotations=args.annotation,
        )
    else:
        project = sf.load_project(root=args.project_root)

    project.extract_tiles(
        qc=qc.Otsu(),
        tile_px=args.tile_px,
        tile_um=tile_um,
        enable_downsample=True,
        normalizer=args.normalizer,
        num_threads=args.num_threads,
    )
    dataset = project.dataset(tile_px=args.tile_px, tile_um=tile_um)

    encoder = sf.build_feature_extractor(args.encoder, weights=args.weights, resize=True)

    project.generate_feature_bags(encoder, dataset, outdir=args.output_dir)

    geojson_dir = args.geojson_dir if args.geojson_dir else os.path.join(args.output_dir, 'geojson')
    save_patch_geojson(args.output_dir, geojson_dir, args.tile_px)


if __name__ == '__main__':
    main()