Reordered some code

stringart/__main__.py

@@ -1,54 +1,84 @@
 import cProfile
 import pstats
-
 import matplotlib.pyplot as plt
-
 from stringart import StringArtGenerator
+import numpy as np
+from tqdm import tqdm
+import time
+from collections import Counter
+import os
 
 if __name__ == '__main__':
+    R_goal = 1450/2  #mm
+    dx_abs = 1       #mm
+    iteration_vec = [1000]
+    detail_level_vec = [120]
+    overwrite_pattern = False  # set to True to overwrite existing pattern files
+    for detail_level in detail_level_vec:
+        for iterations in iteration_vec:
+            file_name = f'sphere-01'
+            extension = 'png'
+            profiler = cProfile.Profile()
+            profiler.enable()
+            nails = int(4553*2)         
+            # iterations = 1000
+            generator = StringArtGenerator()
+            generator.load_image(f"stringart/demo/input/{file_name}.{extension}")
+            generator.preprocess()
+            generator.set_shape('two_circles')
+            R = generator.get_radius()
+            dx = dx_abs/R_goal*R
+
+            generator.set_seed(40)
+            generator.set_iterations(iterations)
+            generator.set_weight(detail_level)
+            generator.set_nails(nails)  # 288
+            pattern_file_name = f'{file_name}_{nails}_{iterations}_{detail_level}.csv'
+
+            # if os.path.exists(pattern_file_name) and not overwrite_pattern:
+            #     print(f"Skipping pattern generation for {pattern_file_name}")
+            #     continue
+
+            pattern = generator.generate()
+
+            pattern_array = np.array(pattern)
+            np.savetxt(pattern_file_name, pattern_array, delimiter=',')
+            counter = Counter(map(tuple, pattern_array))
+            most_common_coordinate, count = counter.most_common(1)[0]
+            print(most_common_coordinate, count)
+            lines_x = []
+            lines_y = []
+            for i, j in zip(pattern, pattern[1:]):
+                lines_x.append((i[0], j[0]))
+                lines_y.append((i[1], j[1]))
 
-    profiler = cProfile.Profile()
-    profiler.enable()
-
-    generator = StringArtGenerator()
-    generator.load_image("demo/input/Sample_ML.jpg")
-    generator.preprocess()
-    generator.set_nails(180)  # 288
-    generator.set_seed(42)
-    generator.set_iterations(4000)
-    pattern = generator.generate()
-
-    lines_x = []
-    lines_y = []
-    for i, j in zip(pattern, pattern[1:]):
-        lines_x.append((i[0], j[0]))
-        lines_y.append((i[1], j[1]))
-
-    xmin = 0.
-    ymin = 0.
-    xmax = generator.data.shape[0]
-    ymax = generator.data.shape[1]
-
-    plt.ion()
-    plt.figure(figsize=(8, 8))
-    plt.axis('off')
-    axes = plt.gca()
-    axes.set_xlim([xmin, xmax])
-    axes.set_ylim([ymin, ymax])
-    axes.get_xaxis().set_visible(False)
-    axes.get_yaxis().set_visible(False)
-    axes.set_aspect('equal')
-    plt.draw()
-
-    batchsize = 10
-    for i in range(0, len(lines_x), batchsize):
-        plt.plot(lines_x[i:i+batchsize], lines_y[i:i+batchsize],
-                 linewidth=0.1, color='k')
-        plt.draw()
-        plt.pause(0.000001)
-
-    plt.savefig('demo/result_ml.png', bbox_inches='tight', pad_inches=0)
+            xmin = 0.
+            ymin = 0.
+            xmax = generator.data.shape[0]
+            ymax = generator.data.shape[1]
+
+            # plt.ion()
+            plt.figure(figsize=(8, 8))
+            plt.axis('off')
+            axes = plt.gca()
+            axes.set_xlim([xmin, xmax])
+            axes.set_ylim([ymin, ymax])
+            axes.get_xaxis().set_visible(False)
+            axes.get_yaxis().set_visible(False)
+            axes.set_aspect('equal')
+
+            batchsize = 10
+            for i in tqdm(range(0, len(lines_x), batchsize), total=len(lines_x)//batchsize, desc='Drawing'):
+                plt.plot(lines_x[i:i+batchsize], lines_y[i:i+batchsize],
+                        linewidth=0.1, color='k')
+
+            # plt.pause(0.000001)
+            plt.savefig(f'stringart/demo/weigth_vec/result_{file_name}_{nails}_{iterations}_{detail_level}_{int(time.time())}.png', bbox_inches='tight', pad_inches=0)
+    plt.show()
+    profiler.disable()
+    # stats = pstats.Stats(profiler).sort_stats('cumtime')
+    # stats.print_stats()
 
     profiler.disable()
-    stats = pstats.Stats(profiler).sort_stats('cumtime')
-    stats.print_stats()
+        # stats = pstats.Stats(profiler).sort_stats('cumtime')
+        # stats.print_stats()

stringart/stringart.py

@@ -9,6 +9,7 @@
 import math
 import copy
 import numpy as np
+from tqdm import tqdm
 
 from PIL import Image, ImageOps, ImageFilter, ImageEnhance
 
@@ -35,10 +36,12 @@ def set_weight(self, weight):
     def set_shape(self, shape):
         self.shape = shape
 
-    def set_nails(self, nails):
+    def set_nails(self, nails, dx=1):
         self.nails = nails
         if self.shape == 'circle':
             self.set_nodes_circle()
+        elif self.shape == 'two_circles':
+            self.set_nodes_two_circles(dx)
         elif self.shape == 'rectangle':
             self.set_nodes_rectangle()
 
@@ -88,6 +91,27 @@ def set_nodes_circle(self):
     def get_radius(self):
         return 0.5*np.amax(np.shape(self.data))
 
+    def set_nodes_two_circles(self, dx):
+        """Set's nails evenly along two circles of given diameter"""
+        spacing = (2*math.pi)/self.nails
+
+        steps = range(self.nails)
+
+        radius1 = self.get_radius() + dx
+        radius2 = self.get_radius() - dx
+
+        x = []
+        y = []
+        for i in steps:
+            if i % 2 == 0:
+                x.append(radius1 + radius1*math.cos(i*spacing))
+                y.append(radius1 + radius1*math.sin(i*spacing))
+            else:
+                x.append(radius2 + radius2*math.cos(i*spacing))
+                y.append(radius2 + radius2*math.sin(i*spacing))
+
+        self.nodes = list(zip(x, y))
+
     def load_image(self, path):
         img = Image.open(path)
         self.image = img
@@ -104,17 +128,18 @@ def preprocess(self):
         self.data = np.flipud(np_img).transpose()
 
     def generate(self):
-        self.calculate_paths()
+        # self.calculate_paths()
         delta = 0.0
         pattern = []
         nail = self.seed
         datacopy = copy.deepcopy(self.data)
-        for i in range(self.iterations):
+        for i in tqdm(range(self.iterations), total=self.iterations, desc='Generating pattern'):
             # calculate straight line to all other nodes and calculate
             # 'darkness' from start node
 
             # choose max darkness path
-            darkest_nail, darkest_path = self.choose_darkest_path(nail)
+            # darkest_nail, darkest_path = self.choose_darkest_path(nail)
+            darkest_nail, darkest_path = self.find_darkest_path(nail)
 
             # add chosen node to pattern
             pattern.append(self.nodes[darkest_nail])
@@ -138,27 +163,47 @@ def generate(self):
 
         return pattern
 
-    def choose_darkest_path(self, nail):
-        max_darkness = -1.0
-        for index, rowcol in enumerate(self.paths[nail]):
-            rows = [i[0] for i in rowcol]
-            cols = [i[1] for i in rowcol]
-            darkness = float(np.sum(self.data[rows, cols]))
+    def calculate_darkness(self, path):
+        rows = [i[0] for i in path]
+        cols = [i[1] for i in path]
+        darkness = float(np.sum(self.data[rows, cols]))
+        return darkness
 
+    def find_darkest_path(self, nail_in):
+        start = self.nodes[nail_in]
+        max_darkness = -1.0
+        for index, end in enumerate(self.nodes):
+            path = self.bresenham_path(start, end)
+            darkness = self.calculate_darkness(path)
             if darkness > max_darkness:
-                darkest_path = np.zeros(np.shape(self.data))
-                darkest_path[rows,cols] = 1.0
                 darkest_nail = index
+                darkest_path = np.zeros(np.shape(self.data))
+                rows = [i[0] for i in path]
+                cols = [i[1] for i in path]
+                darkest_path[rows, cols] = 1.0
                 max_darkness = darkness
-
         return darkest_nail, darkest_path
 
-    def calculate_paths(self):
-        for nail, anode in enumerate(self.nodes):
-            self.paths.append([])
-            for node in self.nodes:
-                path = self.bresenham_path(anode, node)
-                self.paths[nail].append(path)
+    # def choose_darkest_path(self, nail):
+    #     max_darkness = -1.0
+    #     for index, rowcol in enumerate(self.paths[nail]):
+    #         rows = [i[0] for i in rowcol]
+    #         cols = [i[1] for i in rowcol]
+    #         darkness = float(np.sum(self.data[rows, cols]))
+
+    #         if darkness > max_darkness:
+    #             darkest_path = np.zeros(np.shape(self.data))
+    #             darkest_path[rows,cols] = 1.0
+    #             darkest_nail = index
+    #             max_darkness = darkness
+    #     return darkest_nail, darkest_path
+
+    # def calculate_paths(self):
+    #     for nail, anode in tqdm(enumerate(self.nodes), total=len(self.nodes), desc='Calculate paths'):
+    #         self.paths.append([])
+    #         for node in self.nodes:
+    #             path = self.bresenham_path(anode, node)
+    #             self.paths[nail].append(path)
 
     def bresenham_path(self, start, end):
         """Bresenham's Line Algorithm