add support for custom PSF

Phantom_generate/synthetic_data.py

@@ -2,7 +2,9 @@ import os
 import cv2
 import numpy as np
 import random
-from helpers.output import get_save_paths
+
+from helpers.image import show_image
+from helpers.output import get_save_paths,save_array
 """
 This script generates synthetic images with $line_count$ lines of thickness $line_thickness$ and color 
 $line_color$. The images are of size $output_width$ and $output_height$. There is also a blurring function. The 
@@ -14,29 +16,38 @@ the blurred images. N is the run number.
 output_path = "synthetic-data/"
 output_height = 2192
 output_width = 3840
-generation_number = 150  # number of files to output
+generation_number = 5  # number of files to output
 
 line_thickness = (4,15)*255
-line_color = (0,0.05)*255
+
+line_color = (0.98,1.0)
 
 line_count = 20
 square_count = 20
 
-blur_kernel = "gaussian"
+blur_kernel = "custom"
+
+#Options for blur_kernel = "gaussian"
 gaussian_kernel_size = 20
 gaussian_kernel_std_x = 8
 gaussian_kernel_std_y = 16
 
+# Options for blur_kernel = "custom"
+kernel_path = "avg_psf.npy"
 
-def gen_image(height, width, thickness, color, line_count , square_count):
-    image = np.ones((height, width), dtype=np.uint8)*255
+output_png = True
+output_npz = True
 
+interactive = False
+def gen_image(height, width, thickness, color, line_count , square_count):
+    image = np.zeros((height, width), dtype=np.float64)
     for i in range(line_count):
         thickness_i = random.randint(thickness[0],thickness[1])
         color_i = random.uniform(color[0],color[1])
         line_start = (random.randint(0, width-51), random.randint(0, height-51))
         line_end = (random.randint(0, width), random.randint(0, height))
         cv2.line(image, line_start, line_end, color_i, thickness_i)
+
     for j in range(square_count):
         thickness_i = random.randint(thickness[0],thickness[1])
         color_i = random.uniform(color[0],color[1])
@@ -65,9 +76,26 @@ gt_path,blur_path = get_save_paths(output_path)
 if blur_kernel == "gaussian":
     kernel = generate_gaussian_kernel(gaussian_kernel_size, gaussian_kernel_std_x, gaussian_kernel_std_y)
 
+if blur_kernel == "custom":
+    kernel = np.load(kernel_path)
+
+
 for i in range(generation_number):
     gt = gen_image(output_height, output_width, line_thickness, line_color, line_count,square_count)
     blured = cv2.filter2D(gt, -1, kernel)
-    cv2.imwrite(gt_path + str(i) + ".png", gt)
-    cv2.imwrite(blur_path + str(i) + ".png", blured)
+
+    if interactive:
+        show_image(gt, "gt")
+        show_image(blured, "blured")
+        if cv2.waitKey() == ord('c'):
+            print("exiting")
+            cv2.destroyAllWindows()
+            exit(0)
+
+    if output_png:
+        cv2.imwrite(gt_path + str(i) + ".png", gt*255)
+        cv2.imwrite(blur_path + str(i) + ".png", blured*255)
+    if output_npz:
+        save_array(gt_path+str(i)+".png",gt)
+        save_array(blur_path+str(i)+".png",blur_path)
     print(str(i) + " / " + str(generation_number))
\ No newline at end of file