added pre and post processing scripts for testing on Milton

Author: MihinP

partinet/DynamicDet

@@ -58,9 +58,24 @@ def main():
     pass
 
 @main.command()
-def preprocess():
-    click.echo("This will preprocess the micrographs.")
-    raise NotImplementedError("Not implemented yet!")
+@click.option("--labels", required=True, help="Path to the labels directory")
+@click.option("--images", required=True, help="Path to the images directory")
+@click.option("--output", required=True, help="Path to the output directory")
+def preprocess(labels, images, output):
+    click.echo("Preprocessing the micrographs...")
+    import partinet.split_train
+    partinet.split_train.main(labels, images, output)
+
+@main.command()
+@click.option("--labels", required=True, help="Path to the labels directory")
+@click.option("--images", required=True, help="Path to the images directory")
+@click.option("--output", required=True, help="Path to the output STAR file")
+@click.option("--conf", default=0.0, help="Minimum confidence threshold from predictions")
+def star(labels, images, output,conf):
+    click.echo("Generating STAR file...")
+    import partinet.star_file
+    partinet.star_file.main(labels,images,output,conf)
+
 
 @main.group()
 def train():

partinet/__init__.py

@@ -0,0 +1,71 @@
+import os
+import shutil
+import argparse
+
+import numpy as np
+from sklearn.model_selection import train_test_split as tts
+
+def main(labels_path, images_path, output_dir):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    os.makedirs(os.path.join(output_dir, "images"))
+    os.makedirs(os.path.join(output_dir, "labels"))
+    os.makedirs(os.path.join(output_dir, "images", "train"))
+    os.makedirs(os.path.join(output_dir, "images", "val"))
+    os.makedirs(os.path.join(output_dir, "labels", "train"))
+    os.makedirs(os.path.join(output_dir, "labels", "val"))
+
+    files = os.listdir(labels_path)
+    # splitting data into training and val
+    train_idx, val_idx = tts(np.arange(0, len(files), 1), shuffle=True)
+
+    for idx, file in enumerate(files):
+        file_name = file[:-4]
+        if idx in train_idx:
+            # copying train images
+            shutil.copy(os.path.join(images_path, file_name+".jpg"), os.path.join(output_dir, "images", "train", file_name+".jpg"))
+
+            # copying train labels
+            shutil.copy(os.path.join(labels_path, file_name+".txt"), os.path.join(output_dir, "labels", "train", file_name+".txt"))
+
+        elif idx in val_idx:
+            # copying val images
+            shutil.copy(os.path.join(images_path, file_name+".jpg"), os.path.join(output_dir, "images", "val", file_name+".jpg"))
+
+            # copying val labels
+            shutil.copy(os.path.join(labels_path, file_name+".txt"), os.path.join(output_dir, "labels", "val", file_name+".txt"))
+
+    # creating val.txt file
+    with open(os.path.join(output_dir, "val.txt"), "w") as f:
+        for file in os.listdir(os.path.join(output_dir, "images", "val")):
+            f.write(str(os.path.join(output_dir, "images", "val", file))+"\n")
+
+    # creating train.txt file
+    with open(os.path.join(output_dir, "train.txt"), "w") as f:
+        for file in os.listdir(os.path.join(output_dir, "images", "train")):
+            f.write(str(os.path.join(output_dir, "images", "train", file))+"\n")
+
+    # creating cryo_training.yaml file
+    to_write = f"""train: {os.path.join(output_dir, "train.txt")}
+val: {os.path.join(output_dir, "val.txt")}
+
+# number of classes
+nc: 1
+
+# class names
+names: [ 'particle' ]"""
+
+    with open(os.path.join(output_dir, "cryo_training.yaml"), "w") as f:
+        f.write(to_write)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Create training data split from images and labels.")
+    parser.add_argument("--labels", required=True, help="Path to the labels directory")
+    parser.add_argument("--images", required=True, help="Path to the images directory")
+    parser.add_argument("--output", required=True, help="Path to the output directory")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args.labels, args.images, args.output)

partinet/split_train.py

@@ -0,0 +1,104 @@
+import math
+import os
+
+import pandas as pd
+import csv
+import cv2
+import argparse
+
+def yolo_to_starfile(yolo_coords, image_width, image_height,diameters):
+    x_center = math.ceil(yolo_coords['x_centre'] * image_width)
+    y_center = math.ceil(yolo_coords['y_centre'] * image_height)
+    width = yolo_coords['width'] * image_width
+    height = yolo_coords['height'] * image_height
+
+    # Calculate diameter
+    diameter = math.ceil(max(width, height))
+
+    diameters.append(diameter)
+
+    # Return x_center, y_center, diameter
+    return x_center, y_center, diameter
+
+
+# Modify the generate_output function to accept DynamicDet DataFrame
+def generate_output(labels, filename, star_writer,diameters,img_width,img_height):
+    for index, row in labels.iterrows():
+        # x_centre = row['X-Coordinate']
+        # y_centre = row['Y-Coordinate']
+        # diameter = row['Diameter']
+
+        x_centre, y_centre, diameter = yolo_to_starfile(row, img_width, img_height,diameters)
+        # print(x_centre, y_centre, diameter)
+
+        star_writer.writerow([filename, x_centre, y_centre, diameter])
+
+def main(labels_path,images_path,star_out_path,conf_thresh):
+    # # Read DynamicDet output from txt file
+    # labels_path = "path/to/labels"
+    # images_path = "path/to/denoised_images"
+    # star_out_path = "output/path/file.star"
+    # conf_thresh = 0.5
+
+    diameters = list()
+
+    with open(star_out_path, "w") as star_file:
+        star_writer = csv.writer(star_file, delimiter=' ')
+        star_writer.writerow([])
+        star_writer.writerow(["data_"])
+        star_writer.writerow([])
+        star_writer.writerow(["loop_"])
+        star_writer.writerow(["_rlnMicrographName", "#1"])
+        star_writer.writerow(["_rlnCoordinateX", "#2"])
+        star_writer.writerow(["_rlnCoordinateY", "#3"])
+        star_writer.writerow(["_rlnDiameter", "#4"])
+        
+        i = 1
+        for image in os.listdir(images_path):
+            filename = image.split("/")[-1][:-4]
+            
+            print(i)
+            
+            i+=1
+            
+            label_file_path = os.path.join(labels_path, str(filename + '.txt'))
+            
+            if not os.path.exists(label_file_path):
+                continue
+                
+            # label_file_path = os.path.join(labels_path, str(filename + '.csv'))
+            # labels = pd.read_csv(label_file_path)
+
+            image = cv2.imread(os.path.join(images_path, image))
+            img_width = image.shape[1]
+            img_height = image.shape[0]
+
+                
+            custom_headers = ['class', 'x_centre', 'y_centre', 'width', 'height', 'conf']
+
+            # Read the CSV file with custom headers
+            labels = pd.read_csv(label_file_path, header=None, names=custom_headers, sep=' ')
+        
+            labels = labels[labels['conf'] > conf_thresh]
+            
+            # print(labels.head())
+
+            # label_file_path = os.path.join(labels_path, str(filename + '.csv'))
+            # labels = pd.read_csv(label_file_path)
+
+            star_filename = str(filename + '.mrc')
+
+            generate_output(labels, star_filename, star_writer,diameters,img_width,img_height)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate STAR file from PartiNet predictions")
+    parser.add_argument("--labels", required=True, help="Path to the labels directory")
+    parser.add_argument("--images", required=True, help="Path to the images directory")
+    parser.add_argument("--output", required=True, help="Path to the output STAR file")
+    parser.add_argument("--conf", required=True, help="Minimum confidence threshold for predictions")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args.labels, args.images, args.output, args.conf)