refactoring

configs/augmentations.json

@@ -11,8 +11,8 @@
       "param_name": "grid",
       "type": "several_ints",
       "subparam_names": [
-        "y",
-        "x"
+        "height",
+        "width"
       ],
       "limits_list": [
         [
@@ -264,8 +264,8 @@
       "param_name": "tile_grid_size",
       "type": "several_ints",
       "subparam_names": [
-        "y",
-        "x"
+        "height",
+        "width"
       ],
       "limits_list": [
         [

src/main.py

@@ -1,73 +1,74 @@
-import numpy as np
-import os
-import cv2
-import json
-
-import albumentations as A
 import streamlit as st
+import albumentations as A
 
 from control import *
+from utils import (
+    load_image,
+    get_images_list,
+    load_augmentations_config,
+    generate_executable_string,
+)
 
 
-def load_image(image_name, path_to_folder="../images"):
-    path_to_image = os.path.join(path_to_folder, image_name)
-    image = cv2.imread(path_to_image)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+def select_image(path_to_images: str = "images"):
+    image_names_list = get_images_list(path_to_images)
+    image_name = st.sidebar.selectbox("Select an image:", image_names_list)
+    image = load_image(image_name, path_to_images)
     return image
 
 
-st.title("Demo of Albumentations transforms")
+def show_transform_control(transform_params: dict):
+    param_values = {}
+    if len(transform_params) == 0:
+        st.sidebar.text(transform_name + " transform has no parameters")
+    else:
+        for param in transform_params:
+            param_values[param["param_name"]] = param2func[param["type"]](**param)
+    return param_values
 
 
-# selecting the image
-path_to_images = "images"
-image_names_list = [
-    x for x in os.listdir(path_to_images) if x[-3:] in ["jpg", "peg", "png"]
-]
-image_name = st.sidebar.selectbox("Select an image:", image_names_list)
-image = load_image(image_name, path_to_images)
+def show_credentials():
+    st.text("")
+    st.text("")
+    st.subheader("Credentials:")
+    st.text("Source: https://github.com/IliaLarchenko/albumentations-demo")
+    st.text(
+        "Albumentations library: https://github.com/albumentations-team/albumentations"
+    )
+    st.text("Image Source: https://www.pexels.com/royalty-free-images/")
 
 
-# selecting the transformation
-path_to_config = "configs/augmentations.json"
-with open(path_to_config, "r") as config_file:
-    augmentations = json.load(config_file)
-transform_name = st.sidebar.selectbox(
-    "Select a transformation:", sorted(list(augmentations.keys()))
-)
-transform_params = augmentations[transform_name]
+def show_docstring(object):
+    st.subheader("Docstring:")
+    st.text(str(object.__doc__))
 
 
-# show the transform options
-if len(transform_params) == 0:
-    st.sidebar.text(transform_name + " transform has no parameters")
-else:
-    for param in transform_params:
-        param["value"] = param2func[param["type"]](**param)
+# main
+st.title("Demo of Albumentations transforms")
 
+# selecting an image
+image = select_image(path_to_images="images")
 
-params_string = ", ".join(
-    [param["param_name"] + "=" + str(param["value"]) for param in transform_params]
-    + ["p=1.0"]
+# load the config
+augmentations = load_augmentations_config("configs/augmentations.json")
+transform_name = st.sidebar.selectbox(
+    "Select a transformation:", sorted(list(augmentations.keys()))
 )
-params_string = "(" + params_string + ")"
 
-st.text(transform_name + params_string)
+# select the params values
+param_values = show_transform_control(augmentations[transform_name])
+executable_string = generate_executable_string(transform_name, param_values)
+
+st.text(executable_string)
 st.text("Press R to update")
-exec("transform = A." + transform_name + params_string)
+exec("transform = A." + executable_string)
+augmented_image = transform(image=image)["image"]
+
 st.image(
-    [image, transform(image=image)["image"]],
+    [image, augmented_image],
     caption=["Original image", "Transformed image"],
     width=320,
 )
 
-st.subheader("Docstring:")
-st.text(str(transform.__doc__))
-
-
-st.text("")
-st.text("")
-st.subheader("Credentials:")
-st.text("Source: https://github.com/IliaLarchenko/albumentations-demo")
-st.text("Albumentations library: https://github.com/albumentations-team/albumentations")
-st.text("Image Source: https://www.pexels.com/royalty-free-images/")
+show_docstring(transform)
+show_credentials()

src/utils.py

@@ -0,0 +1,36 @@
+import cv2
+import os
+import json
+
+import streamlit as st
+
+
+@st.cache
+def load_image(image_name, path_to_folder="../images"):
+    path_to_image = os.path.join(path_to_folder, image_name)
+    image = cv2.imread(path_to_image)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return image
+
+
+@st.cache
+def get_images_list(path_to_folder):
+    image_names_list = [
+        x for x in os.listdir(path_to_folder) if x[-3:] in ["jpg", "peg", "png"]
+    ]
+    return image_names_list
+
+
+@st.cache
+def load_augmentations_config(path_to_config: str = "configs/augmentations.json"):
+    with open(path_to_config, "r") as config_file:
+        augmentations = json.load(config_file)
+    return augmentations
+
+
+def generate_executable_string(transform_name, param_values):
+    params_string = ", ".join(
+        [k + "=" + str(param_values[k]) for k in param_values.keys()] + ["p=1.0"]
+    )
+    params_string = "(" + params_string + ")"
+    return transform_name + params_string