Filter micrographs by particle count

#!/usr/bin/env python

from cryosparc.tools import CryoSPARC
import json
import argparse
import numpy as np

# you would need to enter your information for CryoSPARC here
with open('/u/rposert/instance-info.json', 'r') as f:
    instance_info = json.load(f)
cs = CryoSPARC(**instance_info)
assert cs.test_connection()

def main(args):
    project_name, select_twod_jobname = args.job_string.split('-')
    
    project = cs.find_project(project_name)
    select_twod = project.find_job(select_twod_jobname)
    particles = select_twod.load_output('particles_selected')
    paths, counts = np.unique(particles['location/micrograph_path'], return_counts = True)
    particles_per_path = dict(zip(paths, counts))
    filtered_paths = {
        path: count for path, count in particles_per_path.items() if count >= args.threshold
    }

    # here you could move files around with shutil, delete things, whatever
    with open('good_paths.txt', 'w') as f:
        f.write('\n'.join(list(filtered_paths.keys())))
        f.write('\n')

parser = argparse.ArgumentParser()
parser.add_argument(
    'job_string',
    help = 'String with project and Select 2D Classes job separated by a dash, e.g., P23-J12'
)
parser.add_argument(
    'threshold',
    help = 'Keep micrographs with this many or more particles in them',
    type = int
)

if __name__ == "__main__":
    args = parser.parse_args()
    main(args)