RLN_code/RLN_single_Train.py

@@ -29,16 +29,23 @@ input_crop_shape = (1000,1000,1) #set to None for whole image
 number_of_images_to_load =200 #set to -1 for all images
 epochs = 400
 
+restore_name = "/RLN_Normed_Model_Clipped_full_0.008_20240429-161854/"
+
 ######## Setup Folders #######
 timestr = time.strftime("%Y%m%d-%H%M%S")
 run_name = "/RLN_Normed_Model_Clipped_{}_{}_{}/".format(model_type,learning_rate,timestr)
 
+if restore_name is not None:
+    run_name =restore_name
 train_model_path = base_dir + '/train/model_rl' + run_name
 train_output = base_dir + '/train/output_rl' + run_name
 test_output = base_dir + '/test/output_rl' + run_name
 log_dir = base_dir + "/logs" + run_name
 
+
 if not os.path.exists(train_model_path):
+    if restore_name is not None: #we are in a restore mode
+        raise Exception("Restore path does not exist")
     os.makedirs(train_model_path)
 
 if not os.path.exists(train_output):
@@ -134,6 +141,7 @@ if model_type =="full":
 if model_type == "reduced":
     model = RLN_model_simple(name="test")
 
+
 with tf.device('/CPU:0'): #load dataset to CPU memory
     dataset_original, (mean,var) = images.move_dataset_to_mem(input_dir, ground_truthdir, number_of_images_to_load)
 
@@ -145,7 +153,10 @@ i = tf.Variable(0, trainable=False, dtype=tf.int64)
 
 ckpt = tf.train.Checkpoint(step=i, model=model)
 manager = tf.train.CheckpointManager(ckpt, train_model_path, max_to_keep=20)
-
+if restore_name is not None:
+    logging.info("Restoring checkpoint from {}".format(restore_name))
+    ckpt.restore(manager.latest_checkpoint)
+    logging.info("Checkpoint restored to {}".format(i))
 #### training Loop ####
 
 for epoch in range(0, epochs):

RLN_code/RLN_single_inference.py

@@ -71,7 +71,7 @@ if __name__ == "__main__":
     arg_parser.add_argument("--base_dir", default=base_dir,
                             help="Base Directory for input using the suggested file structure.")
     arg_parser.add_argument("--run_name", default=None, help="Run name as used in the default file structure.")
-
+    arg_parser.add_argument("-I", "--interactive", action="store_true", help="Show the images interactively.")
     arg_parser.add_argument("--checkpoint_dir", default=None, help="Override checkpoint directory.")
     arg_parser.add_argument("--input_dir", default=None, help="Override input directory.")
     arg_parser.add_argument("--output_dir", default=None, help="Override output directory.")
@@ -107,16 +107,18 @@ if __name__ == "__main__":
     a = os.listdir(input_dir)
     a = sorted(a, key=lambda x: int(x.split(".")[0]))
     print("Input images: {}".format(a))
+    model, i = get_model(train_model_path)
 
     for path_i in a:
         logging.debug(path_i)
         x = images.read_file(input_dir + path_i)
-        x = tf.expand_dims(x, 0)
-        model, i = get_model(train_model_path)
+        if len(x.shape) == 2:  # correct from [width, height] to [width,height, channel]
+            x = tf.expand_dims(x, 0)
+
         st = time.time()
         x_i = inference(model, x)
-
         et = time.time()
+
         dt = np.array(et - st)
         np.savez_compressed(test_output + path_i + ".npz", x=x, f_x=x_i, dt=dt)
         et2 = time.time()

RLN_code/TF2ImageHelpers.py

 import os
 
+
 import numpy as np
 import tensorflow as tf
 import logging
@@ -27,6 +28,12 @@ def read_file(filepath):
     if type == "png":
         file = tf.io.read_file(filepath)
         image = tf.io.decode_png(file)
+    if type == "bmp":
+        file = tf.io.read_file(filepath)
+        image = tf.io.decode_bmp(file)
+        image = tf.cast(image, tf.float32)
+        image = image / 255
+
     elif type == "npz":
         image = load_array(filepath)
         image = np.reshape(image, (image.shape[0], image.shape[1],1))
@@ -94,4 +101,3 @@ def randomize(shape, image, gt):
     g = tf.image.stateless_random_flip_up_down(g, seed)
 
     return i, g
-

helper scripts/PSF/Basler/PSF.py

@@ -16,7 +16,7 @@ this script is for actually determining the PSF for each measurement. For each m
 4) It will save the image
 """
 
-delta = 3# box size
+delta = 7# box size
 detail_display = False # display the position of the brightest point on the whole image
 root = ("/home/guney/Project/Richardson-Lucy-Net/helper scripts/data/psf2")
 dirs = os.listdir(root)