diff --git a/CryoREAD/main.py b/CryoREAD/main.py index bc6d87a..dc1ca79 100644 --- a/CryoREAD/main.py +++ b/CryoREAD/main.py @@ -70,13 +70,15 @@ def init_save_path(origin_map_path): # change contour level to 0 and increase all the density from data_processing.map_utils import increase_map_density - cur_map_path = increase_map_density(cur_map_path, os.path.join(save_path, map_name + "_increase.mrc"), params["contour"]) + cur_map_path = increase_map_density(cur_map_path, os.path.join(save_path, map_name + "_increase.mrc"), + params["contour"]) params["contour"] = 0 # segment map to remove those useless regions from data_processing.map_utils import segment_map cur_new_map_path = os.path.join(save_path, map_name + "_segment.mrc") - cur_map_path = segment_map(cur_map_path, cur_new_map_path, contour=0) # save the final prediction prob array space + cur_map_path = segment_map(cur_map_path, cur_new_map_path, + contour=0) # save the final prediction prob array space # do a pre check to notify user errors for contour with mrcfile.open(cur_map_path, permissive=True) as mrc: @@ -131,7 +133,8 @@ def init_save_path(origin_map_path): Gen_MaskProtein_map(chain_prob, cur_map_path, mask_map_path, params["contour"], threshold=0.3) if params["prediction_only"]: - print("Our prediction results are saved in %s with mrc format for visualization check." % save_path_2nd_stage) + print( + "Our prediction results are saved in %s with mrc format for visualization check." % save_path_2nd_stage) exit() # graph based atomic structure modeling gaussian_bandwidth = params["g"] # use 3 @@ -142,7 +145,8 @@ def init_save_path(origin_map_path): graph_save_path = os.path.join(save_path, "graph_atomic_modeling") mkdir(graph_save_path) - Build_Unet_Graph(cur_map_path, chain_predict_path, fasta_path, graph_save_path, gaussian_bandwidth, dcut, rdcut, params) + Build_Unet_Graph(cur_map_path, chain_predict_path, fasta_path, graph_save_path, gaussian_bandwidth, dcut, rdcut, + params) elif params["mode"] == 1: # structure refinement pipeline diff --git a/main.py b/main.py index d5e9376..2a6c4c0 100644 --- a/main.py +++ b/main.py @@ -11,8 +11,8 @@ logger.add("AutoClass3D.log") parser = argparse.ArgumentParser() - parser.add_argument("-F", type=str, help="Input job folder path containing all MRC files", required=True) - parser.add_argument("-G", type=str, help="GPU ID to use for prediction", required=True) + parser.add_argument("-F", nargs="+", type=str, help="List of input mrc files", required=True) + parser.add_argument("-G", type=str, help="GPU ID to use for prediction", required=False, default="") parser.add_argument("-J", type=str, help="Job name / output folder name", required=True) args = parser.parse_args() @@ -21,12 +21,18 @@ CRYOREAD_PATH = "./CryoREAD/main.py" - mrc_files = glob(args.F + "/*.mrc") + # mrc_files = glob(args.F + "/*.mrc") + + mrc_files = args.F + # check if files exists + for mrc_file in mrc_files: + if not os.path.exists(mrc_file): + logger.error("Input mrc file not found: " + mrc_file) + exit(1) OUTDIR = str(Path("./CryoREAD_Predict_Result").absolute() / args.J) os.makedirs(OUTDIR, exist_ok=True) - # print(mrc_files) logger.info("MRC files count: " + str(len(mrc_files))) logger.info("MRC files path:\n" + "\n".join(mrc_files)) @@ -55,7 +61,8 @@ f"--resolution=8.0", f"--output={curr_out_dir}", ] - process = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True) + process = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + universal_newlines=True) while True: output = process.stdout.readline() @@ -76,7 +83,7 @@ except: logger.warning("Failed to calculate real space CC, maybe the map is empty") real_space_cc = 0.0 - + try: x, fsc, cutoff_05, cutoff_0143 = calculate_fsc(seg_map_path, prot_prob_path) except: