feat:fix:style:docs: Improve and Refactor the program

.gitignore

@@ -114,3 +114,5 @@ dmypy.json
 .pyre/
 
 # End of https://www.gitignore.io/api/python
+
+.vscode/

circle_evolution/evolution.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from circle_evolution.species import Specie
-
+from circle_evolution.render import CircleRenderer
 import circle_evolution.fitness as fitness
 
 
@@ -34,8 +34,9 @@ def __init__(self, target, genes=100):
         self.target = target  # Target Image
         self.generation = 1
         self.genes = genes
+        self.renderer = CircleRenderer((self.size[0], self.size[1]), gray=len(self.size) < 3)
 
-        self.specie = Specie(size=self.size, genes=genes)
+        self.specie = Specie(size=self.size, renderer=self.renderer, genes=genes)
 
     def mutate(self, specie):
         """Mutates specie for evolution.
@@ -46,7 +47,7 @@ def mutate(self, specie):
         Returns:
             New Specie class, that has been mutated.
         """
-        new_specie = Specie(size=self.size, genotype=np.array(specie.genotype))
+        new_specie = Specie(size=self.size, renderer=self.renderer, genotype=np.array(specie.genotype))
 
         # Randomization for Evolution
         y = random.randint(0, self.genes - 1)
@@ -75,7 +76,7 @@ def print_progress(self, fit):
         Args:
             fit (float): fitness value of specie.
         """
-        print("GEN {}, FIT {:.8f}".format(self.generation, fit))
+        print(f"\33[2K\rGEN {self.generation}, FIT {fit:.8f}", end="")
 
     def evolve(self, fitness=fitness.MSEFitness, max_generation=100000):
         """Genetic Algorithm for evolution.

circle_evolution/fitness.py

@@ -26,7 +26,7 @@ def __init__(self, target):
         Args:
             target (np.ndarray): target image array.
         """
-        self.target = target
+        self.target = target.astype(np.float32)
 
     def score(self, phenotype):
         """Score a Specie.

circle_evolution/helpers.py

@@ -1,44 +1,43 @@
 """Helper Functions"""
-import os
 
-import cv2
+import numpy as np
+from PIL import Image
 
-import matplotlib.pyplot as plt
 
-
-def load_target_image(image_path, color=True, size=None):
+def load_target_image(image_path, gray=False, height=None, width=None):
     """Loads images from image path.
 
     Loads and converts image to given colorspace for later processing using
     OpenCV. Attempts to resize image if size is provided.
 
     Args:
         image_path (str): path to load the image.
-        color (bool): if true the image is loaded as rgb, if false grayscale.
-            Defaults to true.
-        size (tuple): size of target image as (height, width). If None, then
-            original image dimension is kept.
+        gray (bool): if True the image is loaded as grayscale, if False rgb.
+            Defaults to False.
+        size (tuple or int): size of target image as (width, height), or width
+            and preserve aspect ratio. If None, then original image dimension is
+            kept.
 
     Returns:
         Image loaded from the path as a numpy.ndarray.
-
-    Raises:
-        FileNotFoundError: image_path does not exist.
     """
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image was not found at {image_path}")
+    target = Image.open(image_path)
 
-    if color:
-        target = cv2.imread(image_path, cv2.IMREAD_COLOR)
-        # Switch from bgr to rgb
-        target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)
+    if gray:
+        target = target.convert("L")
     else:
-        target = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+        target = target.convert("RGB")
 
-    if size:
-        # Only resizes image if it is needed!
-        target = cv2.resize(src=target, dsize=size, interpolation=cv2.INTER_AREA)
-    return target
+    if width is not None and height is not None:
+        target = target.resize((width, height))
+    elif height is not None:
+        width = int(height * target.width / target.height)
+        target = target.resize((width, height))
+    elif width is not None:
+        height = int(width * target.height / target.width)
+        target = target.resize((width, height))
+
+    return np.array(target, dtype=np.uint8)
 
 
 def show_image(img_arr):
@@ -47,7 +46,15 @@ def show_image(img_arr):
     Arguments:
         img_arr (numpy.ndarray): image array to be displayed
     """
-    plt.figure()
-    plt.axis("off")
-    plt.imshow(img_arr / 255)
-    plt.show()
+    img = Image.fromarray(img_arr)
+    img.show()
+
+
+def save_image(img_arr, filename):
+    """Save image to disk.
+
+    Arguments:
+        img_arr (numpy.ndarray): image array to be saved
+    """
+    img = Image.fromarray(img_arr)
+    img.save(filename)

circle_evolution/main.py

@@ -2,37 +2,65 @@
 
 import argparse
 
-import numpy as np
-
 from circle_evolution import __version__
 from circle_evolution.evolution import Evolution
 
 import circle_evolution.helpers as helpers
 
 
-SIZE_OPTIONS = {1: (64, 64), 2: (128, 128), 3: (256, 256), 'auto': None}
-
-
 def main():
     """Entrypoint of application"""
     parser = argparse.ArgumentParser(description=f"Circle Evolution CLI v{__version__}")
 
     parser.add_argument("image", type=str, help="Image to be processed")
-    parser.add_argument("--size", choices=SIZE_OPTIONS.keys(), default='auto', help="Dimension of the image")
-    parser.add_argument("--genes", default=256, type=int, help="Number of genes")
-    parser.add_argument("--max-generations", type=int, default=500000)
+    parser.add_argument("-w", "--width", type=int, default='auto', help="Width of the image used while training")
+    parser.add_argument("-g", "--genes", type=int, default=128, help="Number of genes / circles")
+    parser.add_argument("-m", "--max-generations", type=int, default=50000)
+    parser.add_argument("-c", "--checkpoint", type=str, default=None)
+    parser.add_argument("-l", "--load-checkpoint", type=str, default=None)
+    parser.add_argument("-o", "--output-width", type=int, default=None)
     args = parser.parse_args()
 
-    target = helpers.load_target_image(args.image, size=SIZE_OPTIONS[args.size])
+    target = helpers.load_target_image(args.image, size=args.width)
+    print(f"Image loaded at resolution: {target.shape[1]}x{target.shape[0]}")
 
     evolution = Evolution(target, genes=args.genes)
-    evolution.evolve(max_generation=args.max_generations)
+    print(f"Using GPU '{evolution.renderer.gpu_name}'")
+
+    # TODO add loading and saving checkpoint logic inside of evolution class
+    starting_generation = 0
+    if args.load_checkpoint is not None:
+        evolution.specie.load_checkpoint(args.load_checkpoint)
+        hyphen_index = args.load_checkpoint.rfind("-")
+        dot_index = args.load_checkpoint.rfind(".")
+        if hyphen_index != -1 and dot_index != -1:
+            try:
+                starting_generation = int(args.load_checkpoint[hyphen_index + 1:dot_index])
+            except ValueError:
+                starting_generation = 0
+
+    if args.max_generations > 0:
+        try:
+            evolution.evolve(max_generation=args.max_generations)
+        except KeyboardInterrupt:
+            print("\nEvolution interrupted by user.")
+        else:
+            print("")
 
-    evolution.specie.render()
-    helpers.show_image(evolution.specie.phenotype)
+    # TODO add logic for saving specie to different size than training size
+    if args.output_width is not None:
+        target = helpers.load_target_image(args.image, size=args.output_width)
+        out_evolution = Evolution(target, genes=args.genes)
+        out_evolution.specie.genotype = evolution.specie.genotype
+        out_evolution.specie.render()
+        helpers.show_image(out_evolution.specie.phenotype)
+    else:
+        evolution.specie.render()
+        helpers.show_image(evolution.specie.phenotype)
 
-    output_path_checkpoint = "checkpoint-{}.txt".format(evolution.generation)
-    np.savetxt(output_path_checkpoint, evolution.specie.genotype)
+    if args.checkpoint:
+        output_path_checkpoint = f"{args.checkpoint}-{starting_generation + evolution.generation}.txt"
+        evolution.specie.save_checkpoint(output_path_checkpoint)
 
 
 if __name__ == "__main__":

circle_evolution/render/__init__.py

@@ -0,0 +1,86 @@
+import moderngl
+import numpy as np
+
+from pathlib import Path
+
+
+class CircleRenderer:
+    def __init__(self, size: tuple[int, int], gray=False) -> None:
+        """Render circles using opengl.
+
+        Args:
+            size (tuple): tuple containing height and width of generated image
+                (h, w).
+            gray (bool): whether to render image as grayscale or rgb. Defaults
+                to False (rgb).
+        """
+        self.gray = gray
+
+        # Initialize ModernGL context
+        self._ctx = moderngl.create_context(standalone=True, require=400)
+
+        self.gpu_name = self._ctx.info['GL_RENDERER']
+
+        current_folder = Path(__file__).parent.resolve()
+        with open(current_folder / "base.vert", 'r') as f:
+            vertex_shader = f.read()
+        with open(current_folder / "circles.frag", 'r') as f:
+            fragment_shader = f.read()
+
+        # Compile the shaders
+        self._prog = self._ctx.program(
+            vertex_shader=vertex_shader,
+            fragment_shader=fragment_shader,
+        )
+
+        # Set the image width and height as a vec2 uniform
+        self._height, self._width = size
+        self._prog['iResolution'] = (self._width, self._height)
+
+        # Create a fullscreen quad VAO
+        quad_vertices = np.array([-1.0, -1.0, 1.0, -1.0, -1.0, 1.0, 1.0, 1.0], dtype='f4')
+        self._vbo = self._ctx.buffer(quad_vertices)
+        self._vao = self._ctx.simple_vertex_array(self._prog, self._vbo, 'in_vert')
+
+        # Create a framebuffer to render into
+        num_dim = 1 if self.gray else 3
+        self._fbo = self._ctx.framebuffer(
+            color_attachments=[self._ctx.texture((self._width, self._height), num_dim)],
+        )
+
+        # Use the framebuffer for rendering
+        self._fbo.use()
+
+        # Clear the framebuffer
+        self._ctx.clear()
+
+    def __del__(self):
+        self._fbo.release()
+        self._vao.release()
+        self._vbo.release()
+        self._prog.release()
+        self._ctx.release()
+
+    def render(self, count, pos, radii, colors):
+        self._prog['circleCount'] = count
+        self._prog['pos'] = pos
+        self._prog['radii'] = radii
+        self._prog['colors'] = colors
+
+        # Render the fullscreen quad
+        self._vao.render(moderngl.TRIANGLE_STRIP)
+
+        # Read the rendered image from the framebuffer
+        # And convert the image to a numpy array
+        if self.gray:
+            pixels = self._fbo.read(components=1, dtype='f1')
+
+            image_data = np.frombuffer(pixels, dtype=np.uint8)
+            image_data = image_data.reshape((self._height, self._width))
+        else:
+            pixels = self._fbo.read(components=3, dtype='f1')
+
+            image_data = np.frombuffer(pixels, dtype=np.uint8)
+            image_data = image_data.reshape((self._height, self._width, 3))
+
+        return image_data

circle_evolution/render/base.vert

@@ -0,0 +1,9 @@
+#version 400
+
+in vec2 in_vert;
+out vec2 fragCoord;
+
+void main() {
+    fragCoord = (in_vert + 1.0) * 0.5;
+    gl_Position = vec4(in_vert, 0.0, 1.0);
+}

circle_evolution/render/circles.frag

@@ -0,0 +1,32 @@
+#version 400
+
+vec4 circle(vec2 uv, vec2 pos, float rad, vec4 color)
+{
+    float d = length(pos - uv) - rad;
+    float t = clamp(d, 0.0, color.a);
+    return vec4(color.xyz, color.a - t);
+}
+
+uniform int circleCount;
+uniform vec2 pos[256];
+uniform float radii[256];
+uniform vec4 colors[256];
+
+out vec4 fragColor;
+in vec2 fragCoord;
+uniform vec2 iResolution;
+
+void main()
+{
+    vec2 uv = fragCoord.xy * iResolution;
+
+    vec4 dest = vec4(0, 0, 0, 1); // RGB values
+
+    for (int i = 0; i < circleCount; i++)
+    {
+        vec4 circ = circle(uv, pos[i], radii[i], colors[i]);
+        dest = mix(dest, circ, circ.a);
+    }
+
+    fragColor = dest;
+}