RLN_code/RLN_single_inference.py

+import numpy
 import numpy as np
 import tensorflow as tf
-
+from matplotlib import pyplot as plt
+from matplotlib import cm
 from RLN_single_Model import RLN_model
 
 
@@ -46,12 +48,38 @@ def inference(model, input):
     return output
 
 
+def preprocess(x, min_p=5, max_p=95, invert=True):
+    min = np.percentile(x.numpy(), min_p)
+    x = x - min
+
+    x = tf.clip_by_value(x, 0, 1)
+    max = np.percentile(x.numpy(), max_p)
+    x = x / max
+    x = tf.clip_by_value(x, 0, 1)
+    if invert:
+        x = 1 - x
+    return x
+
+
+def postprocess(x, min_p=5, max_p=95, invert=True):
+    min = np.percentile(x, min_p)
+    x = x - min
+
+    x = tf.clip_by_value(x, 0, 1)
+    max = np.percentile(x, max_p)
+    x = x / max
+    x = tf.clip_by_value(x, 0, 1)
+    if invert:
+        x = 1 - x
+    return x
+
+
 if __name__ == "__main__":
     import logging
     import argparse
     import os
     import TF2ImageHelpers as images
-    import time
+    from helpers.output import get_last_run_number
 
     logging.basicConfig(filemode="a", encoding='utf-8', level=logging.DEBUG)
 
@@ -68,18 +96,32 @@ if __name__ == "__main__":
                                          description="Used to run Inference on input pictures. Should support .bmp, "
                                                      ".png, .npz file formats.")
 
+    #file controls
     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("-C", "--create_run_dir", action="store_true", help="Save output in a separate folder")
     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.")
-
+    # processing controls
+    arg_parser.add_argument("--crop",default=None,help="Set crop region. Should be a string of format 'x,y,w,h' x,y being left corner coordinates.")
+    arg_parser.add_argument("--iter", default=1, help="How many iterations to run")
+    arg_parser.add_argument("--tile_size", default=0, help="Set the size of sub batches for inference. 0 runs "
+                                                           "inference on the entire region of interest as a single "
+                                                           "image.")
+    arg_parser.add_argument("--postprocess", default="none", help="enable post processing: available options are: local, global, both, and none (default)")
+
+    #output controls
+    arg_parser.add_argument("-P","--png", action="store_true", help="Save png images in output directory.")
+    arg_parser.add_argument("-H","--heatmap", action="store_true", help="Save heatmap images in output directory.")
+    arg_parser.add_argument("-N","--numpy", action="store_true", help="Save numpy arrays in output directory.")
     args = arg_parser.parse_args()
     base_dir = args.base_dir
+    create_run_dir = args.create_run_dir
+    iterations = int(args.iter)
 
-    if args.run_name is None or args.input_dir is None or args.output_dir is None:
+    if args.run_name is None and args.input_dir is None and args.output_dir is None:
         raise Exception("You must specify either --run_name, or --input_dir, --output_dir, and --checkpoint_dir.")
 
     if args.checkpoint_dir is not None and args.output_dir is not None and args.run_name is not None and args.input_dir is not None:
@@ -100,29 +142,124 @@ if __name__ == "__main__":
 
     print(f"Input: {input_dir} \n Output: {test_output} \n Checkpoint: {train_model_path}")
 
-    if not os.path.exists(train_model_path) or not os.path.exists(
-            test_output) or not os.path.exists(input_dir):
+    if not os.path.exists(train_model_path) or not os.path.exists(input_dir):
         raise Exception("missing locations")
 
+    if not os.path.exists(test_output):
+        print("missing location {} trying to create.".format(test_output))
+        os.makedirs(test_output)
+
+    if create_run_dir:
+        last_run_no = get_last_run_number(test_output) + 1
+        test_output = test_output + f"/output_{last_run_no}/"
+        print(test_output)
+        if not os.path.exists(test_output):
+            print("missing location {} trying to create.".format(test_output))
+            os.makedirs(test_output)
+
+    postprocess_local = False
+    postprocess_global = False
+    if args.postprocess == "both":
+        postprocess_local = True
+        postprocess_global = True
+    elif args.postprocess == "global":
+        postprocess_local = False
+        postprocess_global = True
+    elif args.postprocess == "local":
+        postprocess_local = True
+        postprocess_global = False
+
     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)
-
+    print("Model loaded at {}".format(i.numpy()))
     for path_i in a:
         logging.debug(path_i)
-        x = images.read_file(input_dir + path_i)
+        try:
+            x = images.read_file(input_dir + path_i)
+        except Exception:
+            logging.info("Error reading image {}".format(path_i))
+            continue
         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()
-        i.assign_add(1)
-        logging.info("Run: {} calc time: {} save time: {} memory: {}".format(i.numpy(), dt, et2 - et,
-                                                                             tf.config.experimental.get_memory_info(
-                                                                                 'GPU:0')))
+        shape = x.shape
+        if args.crop is not None:
+            crops = args.crop.split(",")
+            x_crop = int(crops[0])
+            y_crop = int(crops[1])
+            w_crop = int(crops[2])
+            h_crop = int(crops[3])
+            x = x[y_crop:y_crop+h_crop,x_crop:x_crop+w_crop, :]
+
+        x = preprocess(x)
+
+        delta = int(args.tile_size)
+        if delta != 0:
+            overlap = 50
+            cols = x.shape[1] // delta
+            rows = x.shape[0] // delta
+
+            x_in = np.zeros_like(x.numpy())
+            x_out = np.zeros_like(x_in)
+            progbar = tf.keras.utils.Progbar(rows * cols * iterations)
+            for row in range(rows):
+                for col in range(cols):
+                    #check if we are on the boundary of the image, if so don't take an overlap.
+                    overlap_col_min = overlap if col != 0 else 0
+                    overlap_col_max = overlap if col != cols-1 else 0
+                    overlap_row_min = overlap if row != 0 else 0
+                    overlap_row_max = overlap if row != rows-1 else 0
+
+                    # logging.info(f"entering row: {row} col: {col}")
+                    x_loc = x[row * delta - overlap_row_min:(row + 1) * delta + overlap_row_max,
+                            col * delta - overlap_col_min:(col + 1) * delta + overlap_col_max, :]
+                    # x_loc = preporcess(x_loc)
+
+                    inshape  = x_loc.shape
+                    x_in[row * delta:(row + 1) * delta, col * delta:(col + 1) * delta, :] = x_loc.numpy()[overlap_row_min:inshape[0] - overlap_row_max, overlap_col_min:inshape[1] - overlap_col_max, :]
+
+
+
+                    x_out_loc = x_loc
+                    for i in range(iterations):
+                        x_out_loc = inference(model, x_out_loc)[0]
+                        x_out_loc = tf.clip_by_value(x_out_loc, 0, 1)
+                        if postprocess_local:
+                            x_out_loc = postprocess(x_out_loc, 1, 99, False)
+
+                        progbar.add(1)
+                    outshape = x_out_loc.shape
+                    x_out[row * delta:(row + 1) * delta, col * delta:(col + 1) * delta, :] = x_out_loc.numpy()[overlap_row_min:outshape[0]-(overlap_row_max), overlap_col_min:outshape[1]-(overlap_col_max), :]
+        else:
+            x_in = x.numpy()
+            x_out = inference(model, x)[0]
+        logging.info("stats: in: sum: {} mean: {} out: sum: {} mean: {}".format(np.sum(x_in), np.mean(x_in), np.sum(x_out), np.mean(x_out)))
+
+        if postprocess_global:
+            x_out = postprocess(x_out, 4, 99, False)
+
+        if args.heatmap:
+            fig, ax = plt.subplots(figsize=(15,15))
+            c = ax.imshow(x_out,cmap=cm.coolwarm)
+            ax.axis('off')
+            fig.colorbar(c, ax=ax)
+            fig.savefig(test_output + "out_" + path_i+"_heatmap.png")
+
+            fig, ax = plt.subplots(figsize=(15,15))
+            c = ax.imshow(x_in,cmap=cm.coolwarm)
+            ax.axis('off')
+            fig.colorbar(c, ax=ax)
+            fig.savefig(test_output + "in_" + path_i+"_heatmap.png")
+
+        if args.png:
+            out = tf.io.encode_png(tf.cast(x_out * 255, tf.uint8))
+            # print(out)
+            tf.io.write_file(test_output + "out_" + path_i + ".png", out)
+
+            in_s = tf.io.encode_png(tf.cast(x_in * 255, tf.uint8))
+            tf.io.write_file(test_output + "in_" + path_i + ".png", in_s)
+
+        if args.numpy:
+            numpy.savez(test_output + "out_" + path_i + ".npz", out=x_out, ini=x_in)
+    with open(test_output + "/settings.txt", "w") as f:
+        f.write(f" Chunk Size: {delta} x {delta} \n Iterations: {iterations}\n Postprocess: {args.postprocess}")

RLN_code/TF2ImageHelpers.py

@@ -28,11 +28,13 @@ def read_file(filepath):
     if type == "png":
         file = tf.io.read_file(filepath)
         image = tf.io.decode_png(file)
-    if type == "bmp":
+        image = tf.cast(image, tf.float32)
+        image = image/255.
+    elif type == "bmp":
         file = tf.io.read_file(filepath)
         image = tf.io.decode_bmp(file)
         image = tf.cast(image, tf.float32)
-        image = image / 255
+        image = image / 255.
 
     elif type == "npz":
         image = load_array(filepath)

tf2_image_gpu.def

@@ -2,7 +2,7 @@ Bootstrap: docker
 From: tensorflow/tensorflow:2.16.1-gpu
 
 %post
-
+pip install matplotlib
 %environment
 
 %files