improve geometry pool/traffic light/circle

app.py

@@ -13,19 +13,33 @@ from transformers import pipeline
 
 
 @interactive(
-    radius=(0.005, [0., 0.01]),
-    isotropy=(0.5, [0., 1.]),
+    background_color=("green", ["green", "blue", "red"]),
+    border_size=(0.05, [0., 0.3]),
 )
-def gen_color(
-    radius: int = 0,  # discrete slider (int)
-    isotropy: float = 0.,  # continuous slider (float)
-):
-    out = np.zeros((256, 256, 3))  # Initial background set to black
-    out[:, :, 1] = 0.5
-    x, y = np.meshgrid(np.linspace(-1, 1, 256), np.linspace(-1, 1, 256))
-    # ball positions and colors: [(cx, cy), [R, G, B]]
+def generate_background(background_color="green", border_size: float = 0.) -> np.ndarray:
+    out = np.zeros((256, 128, 3))  # Initial background set to black
+    border_int = int(border_size * 256)
+    out[border_int:out.shape[0]-border_int, border_int:out.shape[1]-border_int,
+        ["red", "green", "blue"].index(background_color)] = 0.5
+    return out
+
+
+@interactive(
+    radius=(0.005, [0., 0.01]),  # continuous slider (float)
+    spread=(1., [0., 2.]),  # continuous slider (float)
+    geometric_shape=("snooker", ["snooker", "circle", "traffic light"]),
+)
+def add_circles(
+    background: np.ndarray,
+    radius: float = 0.,
+    spread: float = 1.,
+    geometric_shape: str = "snooker",
+) -> np.ndarray:
+    out = background.copy()
+    x, y = np.meshgrid(
+        np.linspace(-1, 1, out.shape[1]), np.linspace(-1, 1, out.shape[0]))
     balls = [
-        ((0, 0), [0.8, 0.8, 0.8]),  # Cue ball (white)
+        ((0., 0.3), [0.8, 0.8, 0.8]),  # Cue ball (white)
         ((0.0, -0.6), [1, 1, 0]),
         ((-0.15, -0.85), [1, 0, 0]),
         ((0.0, -0.85), [0, 1, 0]),
@@ -33,18 +47,44 @@ def gen_color(
         ((-0.075, -0.725), [1, 1, 0]),
         ((0.075, -0.725), [1, 0, 0]),
     ]
-    if isotropy <= 0.5:
-        # 0.5 -> 1
-        x_scale, y_scale = isotropy, 1
-    else:
-        x_scale, y_scale = 1, 1-isotropy
-    for (cx, cy), color in balls:
-        r = (x_scale * (x - cx) ** 2 + y_scale * (y - cy) ** 2) < radius
-        out[r, :] = color  # Set the color for each ball based on its position
+    circle_clock = [
+        ((np.cos(angle), np.sin(angle)), [1, 1, 0]) for angle in np.linspace(0, 2*np.pi, 12)
+    ]
+    traffic_light = [
+        ((0.0, 0.0), [1, 0.8, 0]),
+        ((0.0, 0.12), [0, 1, 0]),
+        ((0.0, -0.12), [1, 0, 0])
+    ]
 
+    chosen_pattern = {"circle": circle_clock, "snooker": balls,
+                      "traffic light": traffic_light}[geometric_shape]
+    for (cx, cy), color in chosen_pattern:
+        r = (x - spread*cx) ** 2 + (y - spread*cy) ** 2
+        out[r < radius, :] = color
     return out
 
 
+@interactive(add_stick=(False, "Add black rectangle"))
+def add_details(img: np.ndarray, add_stick: bool = False) -> np.ndarray:
+    out = img.copy()
+    x, y = np.meshgrid(
+        np.linspace(-1, 1, out.shape[1]), np.linspace(-1, 1, out.shape[0]))
+    if add_stick:
+        # out[(np.abs(x)+0.5*np.abs(y)) < 0.3] = 0.  # [0.8, 0.8, 0.]
+        mask = (np.abs(x) < 0.1) * (0.75*np.abs(y) < 0.2)
+        out[mask, :] = 0.
+    return out
+
+
+@interactive(
+    noise_level=(0.05, [0., 0.2]),
+    seed=(42, [-1, 100])
+)
+def add_noise(img: np.ndarray, noise_level: float = 0., seed: int = 42):
+    np.random.seed(seed)
+    return (img + np.random.normal(0, noise_level, img.shape)).clip(0., 1.)
+
+
 @interactive(detect=(True, "Enable classification"))
 def apply_classifier(img, context={}, detect=False):
     result = {}
@@ -80,10 +120,13 @@ def display_result(context={}):
 
 
 def tutorial_pipeline():
-    inp = gen_color()
-    apply_classifier(inp)
+    background = generate_background()
+    foreground = add_details(background)
+    foreground = add_circles(foreground)
+    noisy_input = add_noise(foreground)
+    apply_classifier(noisy_input)
     result_curve = display_result()
-    return inp, result_curve
+    return [[background, foreground], [noisy_input, result_curve]]
 
 
 if __name__ == "__main__":