From 09346cb114b2ce6e5d3858430d5e356b8d9a4899 Mon Sep 17 00:00:00 2001
From: Artemy <artegoser@gmail.com>
Date: Tue, 18 Apr 2023 13:16:16 +0300
Subject: [PATCH] refactor: methods

---
 methods/bulk_methods.py     |  9 +++++++++
 methods/stack_methods.py    | 10 ++++++++++
 processing/bulk.py          |  9 ++-------
 processing/stacking.py      | 14 ++++----------
 tabs/Bulk processing/app.py |  5 +++--
 tabs/Stacking/app.py        | 12 +++++++-----
 6 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/methods/bulk_methods.py b/methods/bulk_methods.py
index 75bcd15..5fb12bd 100644
--- a/methods/bulk_methods.py
+++ b/methods/bulk_methods.py
@@ -2,6 +2,8 @@ from PIL import Image, ImageFilter
 from skimage import feature
 import numpy as np
 
+methods = ["Sharpen", "Edge detection", "Canny edge detection"]
+
 
 def edge(file_name):
     image = Image.open(file_name)
@@ -22,3 +24,10 @@ def canny_edge(file_name):
 def sharpen(file_name):
     image = Image.open(file_name)
     return image.filter(ImageFilter.SHARPEN)
+
+
+methods_funcs = {
+    "Sharpen": sharpen,
+    "Edge detection": edge,
+    "Canny edge detection": canny_edge,
+}
diff --git a/methods/stack_methods.py b/methods/stack_methods.py
index 17f7e7d..54526ae 100644
--- a/methods/stack_methods.py
+++ b/methods/stack_methods.py
@@ -2,6 +2,8 @@ from PIL import Image, ImageChops
 from tqdm import tqdm
 import numpy as np
 
+methods = ["Denoise", "StarTracks", "Noise extractor", "Untrack"]
+
 
 def denoise(files):
     images = [np.asarray(Image.open(file)) for file in tqdm(files)]
@@ -36,3 +38,11 @@ def untrack(files):
         image = im3
 
     return image
+
+
+methods_funcs = {
+    "Denoise": denoise,
+    "StarTracks": startracks,
+    "Noise extractor": noise_extractor,
+    "Untrack": untrack,
+}
diff --git a/processing/bulk.py b/processing/bulk.py
index 4cb3657..62f685d 100644
--- a/processing/bulk.py
+++ b/processing/bulk.py
@@ -1,7 +1,7 @@
 import ffmpeg
 import os
 from processing.utils import generate_name, get_date_text, generate_name_with_file_name
-from methods.bulk_methods import edge, sharpen, canny_edge
+from methods.bulk_methods import methods_funcs
 
 
 def images_to_video(
@@ -31,12 +31,7 @@ def video_to_images(video_path, img_ext):
 
 def bulk_processing(directory, out_directory, method):
     date = get_date_text()
-    if method == "Edge detection":
-        run_bulk(edge, directory, out_directory, date)
-    elif method == "Canny edge detection":
-        run_bulk(canny_edge, directory, out_directory, date)
-    elif method == "Sharpen":
-        run_bulk(sharpen, directory, out_directory, date)
+    run_bulk(methods_funcs[method], directory, out_directory, date)
 
 
 def run_bulk(func, directory, out_directory, date):
diff --git a/processing/stacking.py b/processing/stacking.py
index 5c3ed7b..45200ac 100644
--- a/processing/stacking.py
+++ b/processing/stacking.py
@@ -1,4 +1,6 @@
-from methods.stack_methods import denoise, startracks, noise_extractor, untrack
+from methods.stack_methods import (
+    methods_funcs,
+)
 import os
 from processing.utils import generate_name
 
@@ -6,16 +8,8 @@ from processing.utils import generate_name
 def stacking(dir, method):
     files = os.listdir(dir)
     files = list(map(lambda x: os.path.join(dir, x), files))
-    files = list(filter(lambda x: x.endswith(".png"), files))
 
-    if method == "denoise":
-        img = denoise(files)
-    elif method == "startracks":
-        img = startracks(files)
-    elif method == "noise extractor":
-        img = noise_extractor(files)
-    elif method == "untrack":
-        img = untrack(files)
+    img = methods_funcs[method](files)
 
     name = generate_name()
     img.save(name)
diff --git a/tabs/Bulk processing/app.py b/tabs/Bulk processing/app.py
index 065169f..7948d6a 100644
--- a/tabs/Bulk processing/app.py	
+++ b/tabs/Bulk processing/app.py	
@@ -1,5 +1,6 @@
 import gradio as gr
 from processing.bulk import bulk_processing
+from methods.bulk_methods import methods
 
 with gr.Blocks() as app:
     gr.Markdown(
@@ -11,8 +12,8 @@ with gr.Blocks() as app:
                 placeholder="A directory with many images.", lines=1, label="Directory"
             )
             method = gr.Dropdown(
-                choices=["Edge detection", "Canny edge detection", "sharpen"],
-                value="Sharpen",
+                choices=methods,
+                value=methods[0],
                 label="Method",
             )
 
diff --git a/tabs/Stacking/app.py b/tabs/Stacking/app.py
index 48add4a..41178f1 100644
--- a/tabs/Stacking/app.py
+++ b/tabs/Stacking/app.py
@@ -1,18 +1,20 @@
 import gradio as gr
 from processing.stacking import stacking
+from methods.stack_methods import methods
 
 with gr.Blocks() as app:
     gr.Markdown("Stacking images.")
     with gr.Row():
         with gr.Column():
             directory = gr.Text(
-                placeholder="A directory with many images of the same size", lines=1, label="Directory")
-            methods = gr.Dropdown(
-                choices=["denoise", "startracks", "noise extractor", "untrack"], value="denoise", label="Method")
+                placeholder="A directory with many images of the same size",
+                lines=1,
+                label="Directory",
+            )
+            method = gr.Dropdown(choices=methods, value=methods[0], label="Method")
             submit = gr.Button("Submit")
 
         with gr.Column():
             output = gr.Gallery()
 
-        submit.click(stacking, inputs=[
-            directory, methods], outputs=[output])
+        submit.click(stacking, inputs=[directory, method], outputs=[output])