-
Notifications
You must be signed in to change notification settings - Fork 0
/
find_vesicles.py
127 lines (112 loc) · 4.45 KB
/
find_vesicles.py
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
from vesicle_picker import (
preprocess,
generate_masks,
postprocess,
helpers,
external_import,
external_export
)
from tqdm import tqdm
import argparse
# Read in the job parameters file
parser = argparse.ArgumentParser(description='Description of your script')
parser.add_argument(
'parameters',
type=str,
help='Path to .ini file containing the parameters for vesicle picking.'
)
args = parser.parse_args()
parameters_filepath = args.parameters
parameters = helpers.read_config(parameters_filepath)
# Use the csparc_import module to initialize a cryosparc session
cs = external_import.load_cryosparc(parameters.get('csparc_input', 'login'))
# Define the project
project = cs.find_project(parameters.get('csparc_input', 'PID'))
# Pull micrographs object from the output of a curate exposures job
micrographs = external_import.micrographs_from_csparc(
cs=cs,
project_id=parameters.get('csparc_input', 'PID'),
job_id=parameters.get('csparc_input', 'JID'),
job_type=parameters.get('csparc_input', 'type')
)
# Initialize the model
model = generate_masks.initialize_model(
model_weights_path=parameters.get('segmentation', 'model_weights_path'),
model_type=parameters.get('segmentation', 'model_type'),
device=parameters.get('segmentation', 'device')
)
# Loop over all micrographs in the job directory
for micrograph in tqdm(micrographs[:]):
# Extract the image
header, image_fullres = project.download_mrc(
micrograph["micrograph_blob/path"]
)
image_fullres = image_fullres[0]
# Extract the micrograph UID
uid = micrograph['uid']
# Use the preprocess module to get micrograph ready for segmentation.
# This script uses bilateral filtering,
# can be adjusted for Gaussian if desired.
preprocessed_micrograph = preprocess.preprocess_micrograph(
image_fullres,
downsample=parameters.getint('general', 'downsample'),
lowpass_mode=parameters.get('preprocessing', 'lowpass_mode'),
d=parameters.getint('preprocessing', 'd'),
sigmaColor=parameters.getint('preprocessing', 'sigmaColor'),
sigmaSpace=parameters.getint('preprocessing', 'sigmaSpace'))
# Generate masks with user-optimized parameters
masks = generate_masks.generate_masks(
preprocessed_micrograph,
model,
psize=parameters.getfloat('general', 'psize'),
downsample=parameters.getint('general', 'downsample'),
points_per_side=parameters.getint('segmentation', 'points_per_side'),
points_per_batch=parameters.getint('segmentation', 'points_per_batch'),
pred_iou_thresh=parameters.getfloat('segmentation', 'pred_iou_thresh'),
stability_score_thresh=(
parameters.getfloat('segmentation', 'stability_score_thresh')
),
crop_n_layers=parameters.getint('segmentation', 'crop_n_layers'),
crop_n_points_downscale_factor=(
parameters.getint('segmentation', 'crop_n_points_downscale_factor')
),
crop_nms_thresh=parameters.getfloat('segmentation', 'crop_nms_thresh'),
min_mask_region_area=(
parameters.getint('segmentation', 'min_mask_region_area')
)
)
if len(masks) == 0:
continue
# Use the postprocess module to compute statistics
# on the vesicles for downstream filtering
postprocessed_masks = postprocess.postprocess_masks(
masks,
eval(parameters.get('postprocessing', 'functions')),
preprocessed_micrograph
)
# Modify the ellipse fitting key-value pairs to convert to Angstrom
for mask in postprocessed_masks:
if 'average_radius' in mask:
psize = parameters.getfloat('general', 'psize')
downsample = parameters.getint('general', 'downsample')
mask['average_radius_A'] = (
mask['average_radius'] * psize * downsample
)
mask['semi_minor_A'] = (
mask['semi_minor'] * psize * downsample
)
mask['semi_major_A'] = (
mask['semi_major'] * psize * downsample
)
# Construct the filepath of the masks to save to disk
masks_filename = (
parameters.get('output', 'directory')+str(uid)+"_vesicles.pkl"
)
# Save the vesicles to the output directory
# on a micrograph-by-micrograph basis
external_export.export_masks_to_disk(
postprocessed_masks,
masks_filename,
micrograph_uid=uid,
compression='uint64'
)