genetic.py

genetic.py

import cv2
import numpy as np
from matplotlib import pyplot as plt

IMAGE = "circle.jpg"
#IMAGE = "lisa.png"
#IMAGE = "melita.png"
INITIAL_POPULATION = 500
EPOCHS = 128
DNA_LENGTH = 64
MUTATE_RATE = 100
RADIUS_MAX = 50
POPULATION_HEALTH_DIVISOR = 5

class Pop:
   def __init__(self, shape):
      self.dna = [] 
      self.shape = shape

   def __str__(self):
      print(f"{self}, len is: #{len(self.dna)}")

   def init_dna(self):
      raise NotImplementedError()

   def generate_image_from_dna(self):
      raise NotImplementedError()

   def mutate_dna(self, child, i):
      pass

   def error_function(self, ref_image):
      # Float value of the difference between the two images
      return (np.sum(np.abs(self.generate_image_from_dna().astype(float) - ref_image.astype(float))), 0)

   def breed(self, other):
      child = self.__class__(self.shape)
      for i in range(len(self.dna)):
         if np.random.randint(0,2) == 0:
            child.dna.append(self.dna[i])
         else:
            child.dna.append(other.dna[i])
      # mutate
      for i in range(len(child.dna)):
         if np.random.randint(0,MUTATE_RATE) == 0:
            self.mutate_dna(child, i)
      return child


class PopCircle(Pop):
   def __init__(self, shape):
      super().__init__(shape)

   def init_dna(self):
      # Random circles, store in DNA
      for i in range(DNA_LENGTH):
         center = (np.random.randint(0,self.shape[1]), np.random.randint(0,self.shape[0])) 
         radius = np.random.randint(1,RADIUS_MAX)
         color = (np.random.randint(0,255),np.random.randint(0,255),np.random.randint(0,255))
         self.dna.append((center,radius,color))

   def generate_image_from_dna(self):
      img = np.zeros((self.shape[0],self.shape[1],3), np.uint8)
      for i in range(len(self.dna)):
         cv2.circle(img,(self.dna[i][0][0],self.dna[i][0][1]),self.dna[i][1],self.dna[i][2],-1)
      return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

   def mutate_dna(self, child, i):
      child.dna[i] = (child.dna[i][0], np.random.randint(1,RADIUS_MAX), child.dna[i][2])
      center = (np.random.randint(0,self.shape[1]), np.random.randint(0,self.shape[0])) 
      child.dna[i] = (center, child.dna[i][1], child.dna[i][2])
      
def get_top_n(pop, ref_image, n):
   # Sort by error
   pop.sort(key=lambda x: x.error_function(ref_image)[0])
   return pop[:n]    
      
ref_image = cv2.imread(IMAGE)
ref_image = cv2.cvtColor(ref_image, cv2.COLOR_BGR2GRAY)
population = []
error_values = []

# Generate initial population
for n in range(0, INITIAL_POPULATION):
    #pop = PopPolygon(ref_image.shape)
    pop = PopCircle(ref_image.shape)
    pop.init_dna()
    population.append(pop)
   
first_best = get_top_n(population, ref_image, 1)[0]

for i in range(0, EPOCHS):
    print("Generation: " + str(i))
    top_list = get_top_n(population, ref_image, len(population))
    best_error = top_list[0].error_function(ref_image)[0]
    error_values.append(best_error)
    print("Best: " + str(best_error))

    # Breed 1000 new population based on best
    population_new = []
    for n in range(0, INITIAL_POPULATION):
        try: 
            parent1 = top_list[np.random.randint(0,INITIAL_POPULATION/POPULATION_HEALTH_DIVISOR)]
            parent2 = top_list[np.random.randint(0,INITIAL_POPULATION/POPULATION_HEALTH_DIVISOR)]
            child = parent1.breed(parent2)
            population_new.append(child)
        except Exception as e:
            print(e)
            print(str(parent1))
            print(str(parent2))
            raise e

    population = population_new

top_list = sorted(population, key=lambda x: x.error_function(ref_image)[0])
best = top_list[0]