gp.py

import math
import string
import random
import time
from PIL import Image
from PIL import ImageDraw
from operator import attrgetter

POSITION_MULTIPLIER = 1
RADIUS_MULTIPLIER = 1
two_pass = False

source = "./src.jpg"

class Graph:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.mse = []
        self.max = 755**2
        self.image = Image.new("RGB", (self.width,self.height), (255,255,255))

    def draw(self):
        self.image = Image.new("RGB", (self.width,self.height),(255,255,255))
        self.max = max(self.bot) + 20
        draw = ImageDraw.Draw(self.image)
        for i in range(len(self.mse)):
            if i > 1:
                x1 = int(float(self.width)/(len(self.mse)-1)*(i))
                x2 = int(float(self.width)/(len(self.mse)-1)*(i-1))
                y1 = self.height - (int(float(self.height) / self.max * self.mse[i]))
                y2 = self.height - (int(float(self.height) / self.max * self.mse[i-1]))
                draw.line([(x1,y1),(x2,y2)],(0,0,0))

        self.image.save("/tmp/graph.png")
        return


class Stroke:
    def __init__(self, x, y, r, c):
        self.x = x
        self.y = y
        self.r = r
        self.c = c

    def soft_mutate(self, magnitude):
        if random.randint(0,1) == 1:
            self.x += int(random.uniform( -magnitude, magnitude) * POSITION_MULTIPLIER)
            self.y += int(random.uniform( -magnitude, magnitude) * POSITION_MULTIPLIER)
            if self.x < 0:
                self.x = 0
            if self.x > art.width - 1:
                self.x = art.width -1
            if self.y < 0:
                self.y = 0
            if self.y > art.height - 1:
                self.y = art.height - 1
            self.c = art.getpixel((self.x, self.y))
        else:
            self.r += int(random.uniform( -magnitude, magnitude) * RADIUS_MULTIPLIER)
            if self.r < MIN_RADIUS:
                self.r = MIN_RADIUS
            if self.r > MAX_RADIUS:
                self.r = MAX_RADIUS
        return


class Painting:
    def __init__(self, strokes, image):
        self.strokes = strokes
        self.image = image
        self.error = 90000

    def mean_square_error( self ):
        ar,ag,ab = self.image.split()
        br,bg,bb = art.split()
        errors = []
        for y in range(art.height):
            for x in range(art.width):
                if y > 0 and y < art.height and x > 0 and x < art.width:
                    errors.append(( ar.getpixel((x,y)) - br.getpixel((x,y)) ) ** 2)
                    errors.append(( ag.getpixel((x,y)) - bg.getpixel((x,y)) ) ** 2)
                    errors.append(( ab.getpixel((x,y)) - bb.getpixel((x,y)) ) ** 2)
        mean_error = sum(errors)/ ( art.width * art.height )
        return mean_error

    def duplicate( self ):

        strokes = []
        for stroke in self.strokes:
            strokes.append(Stroke(stroke.x, stroke.y, stroke.r, stroke.c))
        return Painting(strokes, Image.new("RGB", self.image.size, (255,255,255)))

    def create_image(self):
        image = Image.new("RGB", self.image.size, (255,255,255))
        draw = ImageDraw.Draw(image)
        for stroke in self.strokes:
            draw_circle( draw, stroke.x, stroke.y, stroke.r, stroke.c)
        self.image = image
        return image


def draw_circle(draw, x, y, r, c):

    draw.ellipse([x-r,y-r,x+r,y+r],c,c)
    return

def scale_painting( painting, factor ):
    image = Image.new("RGB", (painting.width*factor,painting.height*factor),(255,255,255))
    strokes = []
    for stroke in painting.strokes:
        strokes.append(Stroke(stroke.x*factor,stroke.y*factor,stroke.r*factor,stroke.c))
    painting = Painting(strokes)
    painting.create_image()
    painting.image.save("./scaled_painting.png")
    return painting

def paint( source, num_of_strokes, num_of_tries ):
    global art, MIN_RADIUS, MAX_RADIUS
    graph = Graph(600,400)
    art = Image.open(source)
    art.convert("RGB")
    MIN_RADIUS = int(float(2 )/695*art.width)
    if MIN_RADIUS < 1:
        MIN_RADIUS = 1
    MAX_RADIUS = int(float(66)/695*art.width)
    layers = []
    strokes = []
    print("Creating initial generation")
    prev_error = 999999
    painting = Painting(strokes, Image.new( "RGB", art.size, (255, 255,255)))
    old_image = painting.image
    for stroke in range(num_of_strokes):
        print("STROKE #" + str(stroke))
        x = random.randint(0,art.width - 1)
        y = random.randint(0,art.height- 1)
        r = random.randint(MIN_RADIUS,MAX_RADIUS)
        c = art.getpixel((x,y))
        c2 = old_image.getpixel((x,y))
        n = 0
        ar,ab,ag = c
        br,bb,bg = c2
        av = ar+ab+ag
        bv = br+bb+bg
        while (av-bv)**2 < 32:
            n += 1
            x = random.randint(0,art.width - 1)
            y = random.randint(0,art.height- 1)
            c = art.getpixel((x,y))
            c2 = old_image.getpixel((x,y))
            ar,ab,ag = c
            br,bb,bg = c2
            av = ar+ab+ag
            bv = br+bb+bg
        painting.strokes.append(Stroke(x,y,r,c))
        prev_error = painting.mean_square_error(painting.strokes[stroke])
        loop_num = 0
        while loop_num < num_of_tries:
            old_painting = painting.duplicate()
            painting.strokes[stroke].soft_mutate(5)
            painting.create_image()
            error = painting.mean_square_error(painting.strokes[stroke])
            graph.mse.append(error)
            graph.draw()
            print("MSE: " + str(error))
            if error < prev_error:
                loop_num = 0
                painting.create_image()
                prev_error = error
                painting.image.save("./painting.png")
            else:
                painting = old_painting
                loop_num += 1
    if two_pass:
        for stroke in range(len(painting.strokes)):
            loop_num = 0
            while loop_num < num_of_tries:
                old_painting = painting.duplicate()
                painting.strokes[stroke].soft_mutate(5)
                painting.create_image()
                error = painting.mean_square_error(painting.strokes[stroke])
                graph.mse.append(error)
                graph.draw()
                print("MSE: " + str(error))
                if error < prev_error:
                    loop_num = 0
                    painting.create_image()
                    prev_error = error
                    painting.image.save("./painting.png")
                else:
                    painting = old_painting
                    loop_num += 1
    return painting


painting = paint('./src.jpg', 500, 30)

scale_painting(painting, 10)