Use list of files as input

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

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: