|
| 1 | +from pathlib import Path |
| 2 | + |
| 3 | +import click |
| 4 | + |
| 5 | +from .picker import Picker |
| 6 | + |
| 7 | + |
| 8 | +@click.command() |
| 9 | +@click.option( |
| 10 | + "--radius", |
| 11 | + default=80, |
| 12 | + show_default=True, |
| 13 | + help="Particle radius [A].") |
| 14 | +@click.option( |
| 15 | + "--angpix", |
| 16 | + default=None, |
| 17 | + type=float, |
| 18 | + show_default=True, |
| 19 | + help="Override pixelsize in header [A/pix].") |
| 20 | +@click.option( |
| 21 | + "--contamination_binning", |
| 22 | + default=16, |
| 23 | + show_default=True, |
| 24 | + help="Binning used to filter out contaminations.") |
| 25 | +@click.option( |
| 26 | + "--gaussian_cont", |
| 27 | + is_flag=True, |
| 28 | + default=False, |
| 29 | + show_default=True, |
| 30 | + help="Smooth background by subtracting gaussian-filtered tomogram.") |
| 31 | +@click.option( |
| 32 | + "--sigma_cont", |
| 33 | + default=300, |
| 34 | + show_default=True, |
| 35 | + help="If smoothing background, apply gaussian filter with this sigma [A].") |
| 36 | +@click.option( |
| 37 | + "--stdtimes_cont", |
| 38 | + default=2.5, |
| 39 | + show_default=True, |
| 40 | + help="Voxels this many SD below mean are considered contamination.") |
| 41 | +@click.option( |
| 42 | + "--minsize_cont", |
| 43 | + default=50, |
| 44 | + show_default=True, |
| 45 | + help="Keep only contaminations of this size and above [binned px].") |
| 46 | +@click.option( |
| 47 | + "--dilate_cont", |
| 48 | + default=200, |
| 49 | + show_default=True, |
| 50 | + help="Dilate contamination mask by this much [A].") |
| 51 | +@click.option( |
| 52 | + "--radius_times", |
| 53 | + default=4, |
| 54 | + show_default=True, |
| 55 | + help="Picks can be this close together [particle radii].") |
| 56 | +@click.option( |
| 57 | + "--inhibit", |
| 58 | + is_flag=True, |
| 59 | + default=False, |
| 60 | + show_default=True, |
| 61 | + help="Use more elaborate algorithm to get more picks.") |
| 62 | +@click.option( |
| 63 | + "--detection_z", |
| 64 | + default=200, |
| 65 | + show_default=True, |
| 66 | + help="Only consider picks in a central slab of the volume of this extent [px].") |
| 67 | +@click.option( |
| 68 | + "--stdtimes_pick", |
| 69 | + default=1.5, |
| 70 | + show_default=True, |
| 71 | + help="Only consider picks with a central SD this many SD above mean.") |
| 72 | +@click.option( |
| 73 | + "--remove_edge", |
| 74 | + is_flag=True, |
| 75 | + default=False, |
| 76 | + show_default=True, |
| 77 | + help="Only keep particles with a higher foreground than background SD.") |
| 78 | +@click.argument( |
| 79 | + "tomograms", |
| 80 | + nargs=-1, |
| 81 | + type=click.Path(exists=True)) |
| 82 | +@click.argument( |
| 83 | + "output_dir", |
| 84 | + nargs=1, |
| 85 | + type=click.Path(writable=True, dir_okay=True)) |
| 86 | +def cli(radius, |
| 87 | + angpix, |
| 88 | + contamination_binning, |
| 89 | + gaussian_cont, |
| 90 | + sigma_cont, |
| 91 | + stdtimes_cont, |
| 92 | + minsize_cont, |
| 93 | + dilate_cont, |
| 94 | + radius_times, |
| 95 | + inhibit, |
| 96 | + detection_z, |
| 97 | + stdtimes_pick, |
| 98 | + remove_edge, |
| 99 | + tomograms, |
| 100 | + output_dir): |
| 101 | + """Pick particles based on diameter. |
| 102 | +
|
| 103 | + Takes tomograms as input. Writes all outputs to output_dir. |
| 104 | +
|
| 105 | + """ |
| 106 | + for tomo in tomograms: |
| 107 | + |
| 108 | + tomo = Path(tomo) |
| 109 | + |
| 110 | + print("working on") |
| 111 | + |
| 112 | + p = Picker(tomo, |
| 113 | + output_dir, |
| 114 | + radius, |
| 115 | + contamination_binning, |
| 116 | + angpix) |
| 117 | + |
| 118 | + mask = p.getcont(gaussian_cont, |
| 119 | + sigma_cont, |
| 120 | + stdtimes_cont, |
| 121 | + minsize_cont, |
| 122 | + dilate_cont) |
| 123 | + |
| 124 | + print("Contamination Mask written to..") |
| 125 | + |
| 126 | + boxes, particles = p.detect(mask, |
| 127 | + radius_times, |
| 128 | + inhibit, |
| 129 | + detection_z) |
| 130 | + |
| 131 | + print("Initial Picks") |
| 132 | + |
| 133 | + metrics, threshold = p.prefilt(particles, |
| 134 | + stdtimes_pick) |
| 135 | + |
| 136 | + boxs_XYZ = p.filt(boxes, |
| 137 | + metrics, |
| 138 | + threshold, |
| 139 | + remove_edge) |
| 140 | + |
| 141 | + print(f"Wrote {len(boxs_XYZ)} picks to ...") |
0 commit comments