Add download button

app.py

@@ -1,8 +1,13 @@
+import os
+
 import gradio as gr
 import matplotlib.pyplot as plt
 #from fraunhofer_test import generate_diff
 from zern_generator import Zernike
 from fraunhofer_test import Fraunhofer
+from PIL import Image
+import numpy as np
+import shutil
 css = """
 .container {
     # height: 50vh;
@@ -29,19 +34,25 @@ def get_zern_from_file(file):
 
 zernike = Zernike()
 fraunhofer = Fraunhofer(zernike)
-
+curr_diff = None
 theme = gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)
 with gr.Blocks(theme=theme, css=css) as demo:
     with gr.Row():
         zern_sliders = []
-
         with gr.Column(scale=1):
-            # file_inp = gr.File(label="Загрузка коэффициентов из файла", height="2em")
             with gr.Tab(label="Генератор"):
                 
-                defocus_amt = gr.Number(label="Величина дефокуса", value=2, interactive=True)
+                defocus_amt = gr.Number(
+                    label="Величина дефокуса", 
+                    value=2, 
+                    interactive=True
+                )
                 with gr.Row():
-                    num_zernikes = gr.Number(label="Количество коэффициентов Цернике", value=6, interactive=True)
+                    num_zernikes = gr.Number(
+                        label="Количество коэффициентов Цернике", 
+                        value=6, 
+                        interactive=True
+                    )
                     show_zernikes_btn = gr.Button("Показать")
                 sliders_container = gr.Column(elem_classes=["container"])
             
@@ -53,8 +64,14 @@ with gr.Blocks(theme=theme, css=css) as demo:
                     count = min(count, 100)
                     for i in range(count):
                         term_name = f" - {zernike.term[i]}" if i < len(zernike.term) else ""
-                        sld = gr.Slider(key=i, label=f"Zernike {i+1}{term_name}", minimum=-100, interactive=True,
-                                        step=0.01, value=0)
+                        sld = gr.Slider(
+                            key=i, 
+                            label=f"Zernike {i+1}{term_name}", 
+                            minimum=-100, 
+                            interactive=True,
+                            step=0.01, 
+                            value=0
+                        )
     
                         sld.change(update_zerns, inputs=[gr.Number(value=i, visible=False), sld])
                         sliders_container.add_child(sld)
@@ -64,7 +81,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
             with gr.Tab(label="Настройки"):
                 resolution_num = gr.Number(value=256, label="Разрешение")
            
-            # file_inp.upload(get_zern_from_file, inputs=[file_inp], outputs=[*sliders_container.children])
             
         with gr.Column(scale=5):
             with gr.Tab(label="Экран"):
@@ -77,7 +93,9 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 btn_batch = gr.Button("Рассчитать тестовые экраны")
             with gr.Tab(label="Разница у фокуса"):
                 plot_diff = gr.Plot(label="Разница в фокусе")
+
                 btn_diff = gr.Button("Рассчитать разницу в фокусе")
+                file_download = gr.File(label="Скачать")
 
             def on_button_click(resolution):
                 return zernike.generate_zern_wavefront_fig(*[zerns[key] for key in sorted(zerns)])
@@ -85,10 +103,16 @@ with gr.Blocks(theme=theme, css=css) as demo:
             def on_settings_click(resolution):
                 print(resolution)
                 zernike.set_image_params(npix=resolution)
-            
+
+
+
+
             def on_diff_button_click(defocus_amount):
+                shutil.rmtree("temp", ignore_errors=True)
+                os.makedirs("temp")
+                global curr_diff
                 pl, mn, diff = fraunhofer.generate_diff(defocus_amount, *[zerns[key] for key in sorted(zerns)])
-                
+                curr_diff = diff
                 fig, axs = plt.subplots(1, 3, figsize = (20, 10))
                 axs[0].imshow(pl, cmap='grey')
                 axs[1].imshow(mn, cmap='grey')
@@ -99,12 +123,22 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 axs[1].set_title(f"-{defocus_amount}")
                 axs[2].axis('off')
                 axs[2].set_title(f"Difference")
-                return fig
+
+                a = np.min(curr_diff)
+                b = np.max(curr_diff)
+                img_arr = ((curr_diff - a) / (b - a) * 255).astype(np.uint8)
+                img = Image.fromarray(np.stack([img_arr] * 3, axis=-1), mode='RGB')
+                z = [str(zerns[key]) for key in sorted(zerns)]
+
+                filename = f'temp/{"_".join(z[1:])}.png'
+                img.save(filename)
+
+                return fig, filename
 
 
             btn.click(on_button_click, inputs=[resolution_num], outputs=[plot_2d, plot_3d])
             btn_batch.click(zernike.generate_batch, inputs=[resolution_num], outputs=[plot_batch])
             resolution_num.change(on_settings_click, inputs=[resolution_num])
-            btn_diff.click(on_diff_button_click, inputs=[defocus_amt], outputs=[plot_diff])
+            btn_diff.click(on_diff_button_click, inputs=[defocus_amt], outputs=[plot_diff, file_download])
 
 demo.launch()

fraunhofer_test.py

@@ -2,7 +2,9 @@ import matplotlib.pyplot as plt
 import cupy as cp
 import numpy as np
 import math
-from cupyx.scipy.fft import fft2, fftshift, ifftshift
+from numpy.fft import fft2, fftshift, ifftshift
+from cupyx.scipy import fft as fft_gpu
+
 from zern_generator import Zernike
 
 
@@ -10,7 +12,7 @@ def denoise(screen: np.ndarray):
     return np.where(np.abs(screen) < 1e-4, 0, screen)
 
 
-class Fraunhofer:
+class FraunhoferGPU:
     def __init__(self, zernike: Zernike):
         self.focal_distance = 100 * 10 ** -3
         self.distance = 100 * 10 ** -3
@@ -21,17 +23,15 @@ class Fraunhofer:
 
         self.zernike = zernike
 
-    # def _denoise(screen: np.ndarray):
-    #     return np.where(np.abs(screen) < 1e-4, 0, screen)
 
     def get_focal_point(self,input_screen):
         input_screen_gpu = cp.asarray(input_screen)
-        #phase_mul = ((cp.exp(1j*self.k*self.distance) * cp.exp(1j*self.k*(self.zernike.xv**2 + self.zernike.yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance))
-        U_fraun = fft2(ifftshift(input_screen_gpu), overwrite_x=True)
-        #U_fraun = fftshift(U_fraun * phase_mul)
-        U_fraun = fftshift(U_fraun)
-
-
+        xv = cp.asarray(self.zernike.xv)
+        yv = cp.asarray(self.zernike.yv)
+        phase_mul = ((cp.exp(1j*self.k*self.distance) * cp.exp(1j*self.k*(xv**2 + yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance))
+        U_fraun = fft_gpu.fft2(fft_gpu.ifftshift(input_screen_gpu), overwrite_x=True)
+        U_fraun = fftshift(U_fraun * phase_mul)
+        
         I_fraun = cp.abs(U_fraun)**2
         return I_fraun / cp.max(I_fraun)
 
@@ -52,11 +52,108 @@ class Fraunhofer:
         focal_image_plus = denoise(self.get_focal_point(wf_c*cp.exp(f_plus_c * 1j)))
         focal_image_minus = denoise(self.get_focal_point(wf_c*cp.exp(f_minus_c * 1j)))
 
+        
         return (
             cp.asnumpy(focal_image_plus),
             cp.asnumpy(focal_image_minus),
             cp.asnumpy(focal_image_plus - focal_image_minus)
         )
+        
+class Fraunhofer:
+    def __init__(self, zernike: Zernike):
+        self.focal_distance = 100 * 10 ** -3
+        self.distance = 100 * 10 ** -3
+        self.radius = 5 * 10 ** -3
+        self.span = (-(10 ** -2), 10 ** -2)
+        self.wavelength = 500 * 10 ** -9
+        self.k = 2*math.pi/self.wavelength
+
+        self.zernike = zernike
+
+    # def _denoise(screen: np.ndarray):
+    #     return np.where(np.abs(screen) < 1e-4, 0, screen)
+
+    def get_focal_point(self,input_screen):
+        input_screen_gpu = np.asarray(input_screen)
+        phase_mul = ((np.exp(1j*self.k*self.distance) * np.exp(1j*self.k*(self.zernike.xv**2 + self.zernike.yv**2)/(2*self.distance)))/(1j*self.wavelength*self.distance))
+        U_fraun = fft2(ifftshift(input_screen_gpu))
+        U_fraun = fftshift(U_fraun * phase_mul)
+        #U_fraun = fftshift(U_fraun)
+
+
+        I_fraun = np.abs(U_fraun)**2
+        return I_fraun / np.max(I_fraun)
+
+
+    def generate_diff(self, defocus_amount, *zernikes):
+        zernikes = list(zernikes)
+        if len(zernikes) < 5:
+            zernikes += [0] * (5-len(zernikes))
+
+        wf = self.zernike.generate_zern_screen([1], self.radius)
+
+        f_plus = self.zernike.generate_zern_screen(
+            self.zernike.add_defocus(defocus_amount,zernikes),
+            norm_radius=self.radius
+        )
+        f_minus = self.zernike.generate_zern_screen(
+            self.zernike.add_defocus(-defocus_amount,zernikes), 
+            norm_radius=self.radius
+        )
+
+        wf_c = np.asarray(wf)
+        f_plus_c = np.asarray(f_plus)
+        f_minus_c = np.asarray(f_minus)
+        focal_image_plus = denoise(
+            self.get_focal_point(wf_c*np.exp(f_plus_c * 1j))
+        )
+        focal_image_minus = denoise(
+            self.get_focal_point(wf_c*np.exp(f_minus_c * 1j))
+        )
+        diff = focal_image_plus - focal_image_minus
+        normalized_diff = diff / cp.max(diff)
+
+        return (
+            cp.asnumpy(focal_image_plus),
+            cp.asnumpy(focal_image_minus),
+            cp.asnumpy(normalized_diff)
+        )
+
 
 
+if __name__ == "__main__":
+    zernike = Zernike()
+    zernike.set_image_params(npix=2 ** 8)
+    zoom_factor = 20
+    f = Fraunhofer(zernike=zernike)
+    ratio = (f.wavelength*f.focal_distance)/f.radius
+    zernikes = [[1 if i == j else 0 for i in range(0,24)] for j in range(20,24)]
+    print(zernikes)
+    fig, ax = plt.subplots(4,1)
+    for i, z in enumerate(zernikes):
+        print(zernike._osa_to_noll_list(z))
+        wf_a = zernike.generate_zern_screen([1], 0.01/zoom_factor)
+        wf_c = zernike.generate_zern_screen(z, 0.01/zoom_factor)
+        wf_c_nz = zernike.generate_zern_screen(z, 0.01)
+        wf = wf_a * np.exp(wf_c*1j)
+
+
+        focal = f.get_focal_point(wf)
+
+        
+
+        zoom = ratio*zoom_factor
+        #ax[i][0].imshow(np.abs(wf_c_nz))
+
+        ax[i].imshow(np.abs(focal))
+        ax[i].imshow(np.abs(focal))
+    ax[0].set_title("Zernike 5 - DEFOCUS")
+    ax[1].set_title("Zernike 6 - V-PR-ASTIGMATISM")
+    ax[2].set_title("Zernike 7 - V-TREFOIL")
+    ax[3].set_title("Zernike 8 - V-PR_COMA")
+    ax[0].axis('off')
+    ax[1].axis('off')
+    ax[2].axis('off')
+    ax[3].axis('off')
+    plt.show()
 

generate_dataset.py

@@ -1,6 +1,6 @@
 from zern_generator import Zernike
 from PIL import Image
-from fraunhofer_test import Fraunhofer
+from fraunhofer_test import Fraunhofer, FraunhoferGPU
 import cupy as cp
 import matplotlib.pyplot as plt
 import itertools
@@ -12,69 +12,40 @@ import numpy as np
 
 zernike = Zernike()
 zernike.set_image_params(npix=2 ** 8)
-fraunhofer = Fraunhofer(zernike)
+fraunhofer = FraunhoferGPU(zernike)
 
 
-def generate_image(z):
-
-
-    start_time = time()
+def generate_image(z, group_depth):
     arr = fraunhofer.generate_diff(1, 0, *z)
     a = cp.min(arr[2])
     b = cp.max(arr[2])
     img_arr = ((arr[2] - a) / (b - a) * 255).astype(cp.uint8)
     img = Image.fromarray(np.stack([img_arr] * 3, axis=-1), mode='RGB')
-    generation_end_time = time()
-    if True:
-        coeffs = '_'.join([f'{x:.2f}' for x in z[2:]])
-        if not os.path.exists(f'./screens_gpu/{z[0]}/{z[1]}'):
-            try:
-                os.makedirs(f'./screens_gpu/{z[0]}/{z[1]}')
-            except:
-                pass
-        filename = f'./screens_gpu/{z[0]}/{z[1]}/{coeffs}_.png'
-        img.save(filename)
-        save_end_time = time()
-    else:
-        plt.imshow(img)
-        plt.colorbar()
-        plt.show()
-    return generation_end_time - start_time, save_end_time - generation_end_time
-
-
-def generate_dataset(zs):
-    c = 0
-    window_size = 50
-    times = []
-    start_time = time()
-    total_images = len(zs)
+    
+    coeffs = '_'.join([f'{x:.2f}' for x in z[group_depth:]])
+    group = "/".join([str(x) for x in z[:group_depth]])
+    path = f'./screens_tst/{group}'
+    if not os.path.exists(path):
+        try:
+            os.makedirs(path)
+        except:
+            print("Directory already exists. Skipping.")
+    filename = f'{path}/{coeffs}_.png'
+    img.save(filename)
+    
+    
+
+
+def generate_dataset(range, step, kf_cnt):
+    zs = get_zernike_coeffs(range, step, kf_cnt)
     num_workers = 24
     if num_workers is None:
         num_workers = cpu_count()
 
     with Pool(num_workers) as pool:
         results = []
-
         for result in pool.imap_unordered(generate_image, zs):
-            c += 1
-
             results.append(result)
-
-            current_time = time()
-            iteration_time = current_time - start_time
-            times.append(iteration_time)
-            if len(times) > window_size:
-                times.pop(0)
-            avg_iteration_time = sum(times) / len(times)
-
-            # if c % 10 == 0:
-            #     os.system('cls')
-            #     print(f"Generated {c}/{total_images}")
-            #     elapsed_time = current_time - start_time
-            #     print(f"Elapsed time: {elapsed_time:.2f} s")
-            #     remaining_images = total_images - c
-            #     estimated_remaining_time = avg_iteration_time * remaining_images
-            #     print(f"Estimated remaining time: {estimated_remaining_time:.2f} s")
         for result in results:
             result.get()
 
@@ -92,22 +63,33 @@ def generate_dataset_gpu(zs):
     for z in zs:
 
         generate_image(z)
-
-
-if __name__ == '__main__':
-    rng = (-5, 5)
-    step = 0.5
-    points = tuple(np.linspace(*rng, int((rng[1] - rng[0]) / step + 1)))
+        
+def get_zernike_coeffs(range, step, kf_cnt):
+    points = tuple(np.linspace(*range, int((range[1] - range[0]) / step + 1)))
     points_as_tuples = tuple(float(x) for x in points)
+    return np.array(list(set(itertools.product(points_as_tuples, repeat=kf_cnt))))
 
-    zernikes = np.array(list(set(itertools.product(points_as_tuples, repeat=3))))
-
-    # print("Total images: ", len(zernikes))
-    # total_gen_time, total_save_time = 0, 0
-    # for z in zernikes:
-    #     gen_time, save_time = generate_image(z)
-    #     total_save_time += save_time
-    #     total_gen_time += gen_time
-    # print("Average generation time: ", total_gen_time/len(zernikes))
-    # print("Average save time: ", total_save_time/len(zernikes))
-    generate_dataset(zernikes)
+if __name__ == '__main__':
+    # rng = (-5, 5)
+    # step = 0.5
+    # points = tuple(np.linspace(*rng, int((rng[1] - rng[0]) / step + 1)))
+    # points_as_tuples = tuple(float(x) for x in points)
+
+    zernikes = get_zernike_coeffs((-5,5), 0.25, 2)
+    
+    print("Total images: ", len(zernikes))
+    total_gen_time, total_save_time = 0, 0
+    start = time()
+    for i, z in enumerate(zernikes):
+        generate_image(z, 1)
+        if i % 100 == 0:
+            print(f"{i}/{len(zernikes)}")
+    print("Average generation time: ", (time()-start)/len(zernikes))
+    
+    # generate_dataset(zernikes)
+    # last = 1
+    # for i in range(1,6):
+        
+    #     p = get_zernike_coeffs([-5,5], 0.25, i)
+    #     print(f"Количество коэфф-ов: {i} -> Количество комбинаций: {len(p)}")
+    #     last = len(p)

zern_generator.py

@@ -56,10 +56,12 @@ class Zernike:
         degree = math.ceil((-3 + math.sqrt(8 * num_zerns + 1)) / 2)
         zerns = list(zerns) + [0] * int((((degree + 1) * (degree + 2) / 2) - len(zerns)))
         cart = RZern(degree)
-
+        
+        
         cart.make_cart_grid(self.xv / norm_radius, self.yv / norm_radius)
-
-        screen = np.array(cart.eval_grid(np.array(zerns), matrix=True))
+        
+        screen = cart.eval_grid(np.array(zerns), matrix=True)
+        screen = np.array(screen)
         screen[mask == 0] = 0
         screen[np.isnan(screen)] = 0
         return screen
@@ -113,7 +115,6 @@ class Zernike:
         first = (n * (n + 1)) // 2
         second = abs(m)
         third = 0
-        # I wanted to simplify the logic but it seems to work like a colander
         if m > 0 and n % 4 in [0, 1]:
             third = 0
         elif m < 0 and n % 4 in [2, 3]:
@@ -129,11 +130,16 @@ class Zernike:
         return self._mn_to_noll(*self._osa_to_mn(idx))
 
     def _osa_to_noll_list(self, zernikes):
-        indexed_zerns = sorted(map(lambda x: (self._osa_to_noll(x[0]), x[1]), enumerate(zernikes)), key=lambda x: x[0])
+        indexed_zerns = sorted(
+            map(lambda x: (self._osa_to_noll(x[0]), x[1]), enumerate(zernikes)), 
+            key=lambda x: x[0]
+        )
         used_indeces = [x for x, _ in indexed_zerns]
         mn, mx = min(used_indeces), max(used_indeces)
         missing_indeces = set(range(mn, mx + 1)) - set(used_indeces)
-        return [x[1] for x in list(sorted(indexed_zerns + [(x, 0) for x in missing_indeces]))]
+        return [x[1] for x in 
+                list(sorted(indexed_zerns + [(x, 0) for x in missing_indeces]))
+        ]
 
 
 term = ["PISTON", "V-TILT", "H-TILT", "DEFOCUS", "O-PR-ASTIGMATISM", "V-PR-ASTIGMATISM", "V-PR-COMA", "H-PR-COMA",