Advertisement
Guest User

sketch.py

a guest
Oct 13th, 2017
218
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 18.96 KB | None | 0 0
  1.  
  2. # standard library modules
  3. import argparse
  4. import os
  5. import sys
  6. import random
  7. import math
  8. import time
  9. import textwrap
  10. import msvcrt
  11.  
  12. # custom modules
  13. try:
  14.     from PIL import Image, ImageDraw
  15. except ImportError as error:
  16.     error_msg = ('ImportError: {} found in the current working directory or '
  17.                  'any PYTHONPATH location.'.format(error))
  18.     sys.stderr.write('\n'.join(textwrap.wrap(error_msg, 78)) + '\n')
  19.  
  20.     error_msg = ('Please ensure the module file is in the current working '
  21.                  'directory or a PYTHONPATH location before executing the '
  22.                  'program again')
  23.     sys.stderr.write('\n'.join(textwrap.wrap(error_msg, 78)) + '\n')
  24.     sys.stderr.flush()
  25.     time.sleep(0.25)
  26.     print('\nPress any key to quit this program.')
  27.     msvcrt.getch()
  28.     sys.exit()
  29.  
  30.  
  31. def opt_parse():
  32.     """Command-line argument parser"""
  33.  
  34.     description = ("""A program or /g/ divertimento that draws a somewhat artistic
  35.                      rendition of a source image via the use of shapes and a variety
  36.                      of chosen parameters.
  37.                      Every time the program runs a directory is created and around
  38.                      50 images in progress are saved there for the sole reason of
  39.                      amusement and webm making. The directory name is chosen from
  40.                      among parameters values.""")
  41.  
  42.     usage = "%(prog)s [-h] -s [line|circle|square] -l <LENGTH> -t <THICKNESS> -n <ITERS> -i <INFILE> -o <OUTFILE> -a <NNN-NNN> -c [1|2|3] -e <ED ROC>"
  43.     parser = argparse.ArgumentParser(description=description, usage=usage)
  44.  
  45.     help = "specifies the shape to use. Default: line"
  46.     parser.add_argument("-s", "--shape", dest="shape", action="store", metavar="[line|circle|square]", default="line",
  47.                         help=help)
  48.     help = "specifies line length, circle radius or square side length depending what shape is used. Default: 30"
  49.     parser.add_argument("-l", "--length", dest="length", type=int, metavar="<LENGTH>", default=30,
  50.                         help=help)
  51.     help = "specifies line thickness or the square breadth. Default: 1"
  52.     parser.add_argument("-t", "--thickness", dest="thickness", type=int, metavar="<THICKNESS>", default=1,
  53.                         help=help)
  54.     help = "specifies the maximun number of iterations to perform. Default: 50000"
  55.     parser.add_argument("-n", "--iters", dest="iterations", type=int, metavar="<ITERS>", default=50000,
  56.                         help=help)
  57.     help = "specifies the image filename to serve as source"
  58.     parser.add_argument("-i", "--infile", dest="infile", action="store", metavar="<INFILE>", required=True,
  59.                         help=help)
  60.     help = "specifies the image filename to save the final result"
  61.     parser.add_argument("-o", "--outfile", dest="outfile", action="store", metavar="<OUTFILE>", required=True,
  62.                         help="help")
  63.     help = "specifies the range of angles to use while drawing lines. Ex: 010-359 signifies a range between 10 to 359 degrees. Default: 000-359"
  64.     parser.add_argument("-a", "--angle_range", dest="angle_range", action="store", metavar="<NNN-NNN>", default="000-359",
  65.                         help=help)
  66.     help = "specifies how color distribution is sampled. use 1 for weighted whole image sampling, 2 for whole image sampling and 3 for weighted image sampling only from the area to be drawn. Default: 2"
  67.     parser.add_argument("-c", "--color", dest="color", type=int, metavar="[1|2|3]", default=2,
  68.                         help=help)
  69.     help = "specifies the minimum Euclidean distance change of rate to reach before the pregram stops. A value of zero ignores this setting and the program clompletes all iterations. Default: 0.0"
  70.     parser.add_argument("-e", "--ed_roc", dest="ED_RoC", action="store", type=float, metavar="<ED ROC>", default=0.0,
  71.                         help=help)
  72.  
  73.     options = argparse.Namespace()
  74.     options = parser.parse_args(namespace=options)
  75.  
  76.     return options
  77.  
  78.  
  79. start = time.time()
  80.  
  81. # finally parse options
  82. options = opt_parse()
  83.  
  84. # arguments parsing
  85. shape = options.shape
  86. infile = options.infile
  87. iterations = options.iterations
  88. length = options.length
  89. thickness = options.thickness
  90. outfile = options.outfile
  91. angle_range = (int(options.angle_range[:3]), int(options.angle_range[4:]))
  92. color_type = options.color
  93. ED_RoC = options.ED_RoC
  94.  
  95. # directory name and creation
  96. if shape == 'line':
  97.     dirname = "{}-{}-{}-{}-[{}]-{}".format(outfile[:-4],
  98.                                            shape, length, thickness,
  99.                                            angle_range,
  100.                                            color_type)
  101. else:
  102.     dirname = "{}-{}-{}-{}--{}".format(outfile[:-4],
  103.                                        shape, length, thickness,
  104.                                        color_type)
  105.  
  106. if not os.path.exists(dirname):
  107.     os.makedirs(dirname)
  108.  
  109. # open source image and conver
  110. sourceImage = Image.open(infile)
  111. d = sourceImage.convert("RGB")
  112.  
  113. # make temp image
  114. tempImage = Image.new(sourceImage.mode, sourceImage.size, "black")
  115.  
  116.  
  117. def buildBuffer():
  118.     ''' builds data buffers for source and temp images and weighted and
  119.    unweighted color palletes from the original.
  120.    '''
  121.     sourceBuffer = dict()
  122.     tempBuffer = dict()
  123.     colorsBufferWeighted = []
  124.     for x in range(0, sourceImage.size[0]):
  125.         for y in range(0, sourceImage.size[1]):
  126.             sourceBuffer[(x, y)] = sourceImage.getpixel((x, y))
  127.             tempBuffer[(x, y)] = (0, 0, 0)
  128.             colorsBufferWeighted.append(sourceImage.getpixel((x, y)))
  129.     colorsBuffer = list(set(colorsBufferWeighted))
  130.     return sourceBuffer, tempBuffer, colorsBufferWeighted, colorsBuffer
  131.  
  132.  
  133. def templateSquare(length, breadth):
  134.     '''Creates a pixel mask with the required shape.
  135.    '''
  136.     pixels = list()
  137.     d = max(length, breadth)
  138.     tempImage = Image.new('1', (d + 10, d + 10), 0)
  139.     draw = ImageDraw.Draw(tempImage)
  140.     draw.rectangle((0, 0, length, breadth), fill=1, outline=1)
  141.     for x in range(0, tempImage.size[0]):
  142.         for y in range(0, tempImage.size[1]):
  143.             if tempImage.getpixel((x, y)) == 1:
  144.                 pixels.append((x, y))
  145.     tempImage.close()
  146.     return pixels
  147.  
  148.  
  149. def pixelSquare(template, length, breadth):
  150.     '''Creates a pixel field from the mask and a random point then makes
  151.    sure to cull any that lies outside the image's dimensions.
  152.    '''
  153.     pixels = set()
  154.     d = max(length, breadth)
  155.     x = random.randint(0 - d, sourceImage.size[0] + d)
  156.     y = random.randint(0 - d, sourceImage.size[1] + d)
  157.  
  158.     for i, g in template:
  159.         if 0 <= (x + i) < sourceImage.size[0] and 0 <= (y + g) < sourceImage.size[1]:
  160.             pixels.add((x + i, y + g))
  161.     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  162.  
  163.  
  164. def templateCircle(radius):
  165.     '''Creates a pixel mask with the required shape
  166.    '''
  167.     pixels = list()
  168.     d = 2 * radius
  169.     tempImage = Image.new('1', (d + 10, d + 10), 0)
  170.     draw = ImageDraw.Draw(tempImage)
  171.     draw.ellipse((0, 0, d, d), fill=1, outline=1)
  172.     for x in range(0, tempImage.size[0]):
  173.         for y in range(0, tempImage.size[1]):
  174.             if tempImage.getpixel((x, y)) == 1:
  175.                 pixels.append((x, y))
  176.     tempImage.close()
  177.     return pixels
  178.  
  179.  
  180. def pixelCircle(template, radius):
  181.     '''Creates a pixel field from the mask and a random point then makes
  182.    sure to cull any that lies outside the image's dimensions.
  183.    '''
  184.     pixels = set()
  185.  
  186.     x = random.randint(0 - radius, sourceImage.size[0] + radius)
  187.     y = random.randint(0 - radius, sourceImage.size[1] + radius)
  188.  
  189.     for i, g in template:
  190.         if 0 <= (x + i) < sourceImage.size[0] and 0 <= (y + g) < sourceImage.size[1]:
  191.             pixels.add((x + i, y + g))
  192.     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  193.  
  194.  
  195. def templateLine(length, thickness, angle):
  196.     '''Creates a pixel mask with the required shape
  197.    '''
  198.     x2 = int(length * math.cos(math.radians(angle)))
  199.     y2 = int(length * math.sin(math.radians(angle)))
  200.  
  201.     pixels = list()
  202.     tempImage = Image.new('1', (length * 2, length * 2), 0)
  203.     draw = ImageDraw.Draw(tempImage)
  204.     draw.line((length, length, x2 + length, y2 + length), fill=1, width=thickness)
  205.  
  206.     for x in range(0, tempImage.size[0]):
  207.         for y in range(0, tempImage.size[1]):
  208.             if tempImage.getpixel((x, y)) == 1:
  209.                     pixels.append((x - length, y - length))
  210.     tempImage.close()
  211.     return pixels
  212.  
  213.  
  214. def pixelLine_fixedAngle(template, length):
  215.     '''Creates a pixel field from the mask and a random point then makes
  216.    sure to cull any that lies outside the image's dimensions.
  217.    '''
  218.     pixels = set()
  219.  
  220.     x = random.randint(0 - length, sourceImage.size[0] + length)
  221.     y = random.randint(0 - length, sourceImage.size[1] + length)
  222.  
  223.     for i, g in template:
  224.         if 0 <= (x + i) < sourceImage.size[0] and 0 <= (y + g) < sourceImage.size[1]:
  225.             pixels.add((x + i, y + g))
  226.     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  227.  
  228.  
  229. ############## DOES NOT IMPROVE SHIT ###################
  230. # def template2Line(length, thickness):
  231. #     """
  232. #     """
  233. #     template = dict()
  234. #
  235. #     for zeta in range(0, 360):
  236. #         points = list()
  237. #         # Xiaolin Wu's line algorithm
  238. #         x1, y1 = 0, 0
  239. #         # get second Point
  240. #         x2 = int(length * math.cos(math.radians(zeta)))
  241. #         y2 = int(length * math.sin(math.radians(zeta)))
  242. #
  243. #         dx, dy = x2 - x1, y2 - y1
  244. #         is_steep = abs(dx) < abs(dy)
  245. #
  246. #         def p(px, py, is_steep):
  247. #             if not is_steep:
  248. #                 return (px, py)
  249. #             else:
  250. #                 return (py, px)
  251. #
  252. #         if is_steep:
  253. #             x1, y1, x2, y2, dx, dy = y1, x1, y2, x2, dy, dx
  254. #
  255. #         if x2 < x1:
  256. #             x1, x2, y1, y2 = x2, x1, y2, y1
  257. #
  258. #         grad = float(dy) / dx
  259. #         intery = y1 + (1 - (x1 - int(x1))) * grad
  260. #
  261. #         def draw_endpoint(pt):
  262. #             x, y = pt
  263. #             xend = round(x)
  264. #             yend = y + grad * (xend - x)
  265. #             # xgap = 1 - ((x + 0.5) - int((x + 0.5)))
  266. #             px, py = int(xend), int(yend)
  267. #             points.append(p(px, py, is_steep))
  268. #             points.append(p(px, py + 1, is_steep))
  269. #             return px
  270. #
  271. #         xstart = draw_endpoint((x1, y1)) + 1
  272. #         xend = draw_endpoint((x2, y2))
  273. #
  274. #         for x in range(xstart, xend):
  275. #             y = int(intery)
  276. #             points.append(p(x, y, is_steep))
  277. #             points.append(p(x, y + 1, is_steep))
  278. #             intery += grad
  279. #
  280. #         template[zeta] = list(set(points))
  281. #
  282. #     return template
  283. #
  284. #
  285. # def pixel2Line(template, length):
  286. #     '''
  287. #     '''
  288. #     pixels = set()
  289. #
  290. #     x = random.randint(0, sourceImage.size[0] - 1)
  291. #     y = random.randint(0, sourceImage.size[1] - 1)
  292. #
  293. #     for h in range(thickness):
  294. #         for i, g in template:
  295. #             if 0 <= (x + h + i) < sourceImage.size[0] and 0 <= (y + h + g) < sourceImage.size[1]:
  296. #                 pixels.add((x + h + i, y + h + g))
  297. #     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  298.  
  299.  
  300. def pixelLine_one_thickness(length, thickness):
  301.     '''Generates a pixel field from the geometry of a line drawn
  302.    randomly onto the image.
  303.    '''
  304.     pixels = set()
  305.     points = list()
  306.     # get random angle
  307.     zeta = random.randint(angle_range[0], angle_range[1])
  308.  
  309.     # Bresenham's Line Algorithm
  310.     x1, y1 = 0, 0
  311.     # get second Point
  312.     x2 = int(length * math.cos(math.radians(zeta)))
  313.     y2 = int(length * math.sin(math.radians(zeta)))
  314.  
  315.     dx, dy = x2 - x1, y2 - y1
  316.     is_steep = abs(dx) < abs(dy)
  317.  
  318.     # Rotate line
  319.     if is_steep:
  320.         x1, y1, x2, y2 = y1, x1, y2, x2
  321.  
  322.     if x2 < x1:
  323.         x1, x2, y1, y2 = x2, x1, y2, y1
  324.  
  325.     # Recalculate differentials
  326.     dx, dy = x2 - x1, y2 - y1
  327.  
  328.     # Calculate error
  329.     error = int(dx / 2.0)
  330.     ystep = 1 if y1 < y2 else -1
  331.  
  332.     # Iterate over bounding box generating points between start and end
  333.     y = y1
  334.     points = []
  335.     for x in range(x1, x2 + 1):
  336.         coord = (y, x) if is_steep else (x, y)
  337.         points.append(coord)
  338.         error -= abs(dy)
  339.         if error < 0:
  340.             y += ystep
  341.             error += dx
  342.  
  343.     x = random.randint(0, sourceImage.size[0] - 1)
  344.     y = random.randint(0, sourceImage.size[1] - 1)
  345.  
  346.     for h in range(thickness):
  347.         for i, g in points:
  348.             if 0 <= (x + h + i) < sourceImage.size[0] and 0 <= (y + h + g) < sourceImage.size[1]:
  349.                 pixels.add((x + h + i, y + h + g))
  350.     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  351.  
  352.  
  353. def pixelLine(length, thickness):
  354.     '''Generates a pixel list from the geometry of a line drawn
  355.    randomly onto the image.
  356.    '''
  357.     pixels = set()
  358.     points = list()
  359.     # get random angle
  360.     zeta = random.randint(angle_range[0], angle_range[1])
  361.  
  362.     # Xiaolin Wu's line algorithm
  363.     x1, y1 = 0, 0
  364.     # get second Point
  365.     x2 = int(length * math.cos(math.radians(zeta)))
  366.     y2 = int(length * math.sin(math.radians(zeta)))
  367.  
  368.     dx, dy = x2 - x1, y2 - y1
  369.     is_steep = abs(dx) < abs(dy)
  370.  
  371.     def p(px, py, is_steep):
  372.         if not is_steep:
  373.             return (px, py)
  374.         else:
  375.             return (py, px)
  376.  
  377.     if is_steep:
  378.         x1, y1, x2, y2, dx, dy = y1, x1, y2, x2, dy, dx
  379.  
  380.     if x2 < x1:
  381.         x1, x2, y1, y2 = x2, x1, y2, y1
  382.  
  383.     grad = float(dy) / dx
  384.     intery = y1 + (1 - (x1 - int(x1))) * grad
  385.  
  386.     def draw_endpoint(pt):
  387.         x, y = pt
  388.         xend = round(x)
  389.         yend = y + grad * (xend - x)
  390.         # xgap = 1 - ((x + 0.5) - int((x + 0.5)))
  391.         px, py = int(xend), int(yend)
  392.         points.append(p(px, py, is_steep))
  393.         points.append(p(px, py + 1, is_steep))
  394.         return px
  395.  
  396.     xstart = draw_endpoint((x1, y1)) + 1
  397.     xend = draw_endpoint((x2, y2))
  398.  
  399.     for x in range(xstart, xend):
  400.         y = int(intery)
  401.         points.append(p(x, y, is_steep))
  402.         points.append(p(x, y + 1, is_steep))
  403.         intery += grad
  404.  
  405.     x = random.randint(0, sourceImage.size[0] - 1)
  406.     y = random.randint(0, sourceImage.size[1] - 1)
  407.  
  408.     for h in range(thickness):
  409.         for i, g in points:
  410.             if 0 <= (x + h + i) < sourceImage.size[0] and 0 <= (y + h + g) < sourceImage.size[1]:
  411.                 pixels.add((x + h + i, y + h + g))
  412.     return pixels or {(sourceImage.size[0] - 1, sourceImage.size[1] - 1)}
  413.  
  414.  
  415. def colorPixels(pixels):
  416.     '''Builds a color palette from the given pixel field extracting the data
  417.    from the source.
  418.    '''
  419.     colors = []
  420.     for pixel in pixels:
  421.         colors.append(sourceBuffer[pixel])
  422.     return colors
  423.  
  424. # relic
  425. # def lineDraw(pixels, color):
  426. #     '''Draws line'''
  427. #     for pixel in pixels:
  428. #         tempImage.putpixel(pixel, color)
  429.  
  430.  
  431. def BufferDraw(buffer):
  432.     '''Draws buffer to temp image'''
  433.     for item in buffer:
  434.         tempImage.putpixel(item, buffer[item])
  435.  
  436.  
  437. def compare(buffer1=None, buffer2=None, pixels=None, color=None):
  438.     '''Calculates Euclidean distance between two image buffers or the origina
  439.    image buffer and a pixel field representing a shape.'''
  440.     r, g, b = 0, 0, 0
  441.     for pixel in pixels:
  442.         (a, b, c) = sourceBuffer[pixel]
  443.         if color:
  444.             (x, y, z) = color
  445.         else:
  446.             (x, y, z) = tempBuffer[pixel]
  447.         r += (a - x) ** 2
  448.         g += (b - y) ** 2
  449.         b += (c - z) ** 2
  450.     return math.sqrt(r + g + b) / len(pixels)
  451.  
  452.  
  453. sourceBuffer, tempBuffer, colorsBufferWeighted, colorsBuffer = buildBuffer()
  454. # relic form from when I was experimenting mixing these.
  455. # may play with it some more later
  456. colors = {1: [100, 0, 0],
  457.           2: [0, 100, 0],
  458.           3: [0, 0, 100]}
  459.  
  460. # pixel field template to use if any
  461. if shape == 'circle':
  462.     template = templateCircle(length)
  463. if angle_range[0] == angle_range[1] and shape == 'line':
  464.     template = templateLine(length, thickness, random.randint(angle_range[0], angle_range[1]))
  465. if shape == 'square':
  466.     template = templateSquare(length, thickness)
  467.  
  468. # main loop
  469. counter, pre_ed = 0, 0
  470. changes, changes_p, changes_t = (0, 0, 0)
  471.  
  472. start_l = time.time()
  473. for i in range(iterations):
  474.     changes_t += 1
  475.     # generates pixel list representing line
  476.     if thickness == 1 and shape == 'line':
  477.         p = pixelLine_one_thickness(length, 1)
  478.     elif angle_range[0] == angle_range[1] and shape == 'line':
  479.         p = pixelLine_fixedAngle(template, length)
  480.     elif shape == 'line':
  481.         p = pixelLine(length, thickness)
  482.     # generates pixel list representing shape
  483.     elif shape == 'circle':
  484.         p = pixelCircle(template, length)
  485.     elif shape == 'square':
  486.         p = pixelSquare(template, length, thickness)
  487.     # picks a random color from the original image
  488.     color = random.choices([random.choice(colorsBufferWeighted),
  489.                             random.choice(colorsBuffer),
  490.                             random.choice(colorPixels(p))],
  491.                            colors[color_type], k=1)[0]
  492.  
  493.     # calculates Euclidean distances
  494.     a = compare(buffer1=sourceBuffer, buffer2=tempBuffer, pixels=p)
  495.     b = compare(buffer1=sourceBuffer, pixels=p, color=color)
  496.  
  497.     # if shape improves tenp image then draw shape
  498.     if a > b:
  499.         for pixel in p:
  500.             tempBuffer[pixel] = color
  501.         changes += 1
  502.         changes_p += 1
  503.  
  504.     # report progress every 10000 iterations
  505.     if i % 10000 == 0 and i != 0:
  506.         end_l = time.time()
  507.         text = "{:d}/{:d} iter ({:d} accepted, {:d} discarded, ratio {:.02f}), {:02.02f}% {:.0f}s."
  508.         print(text.format(i, iterations,
  509.                           changes_p, changes_t - changes_p, float(changes_p) / (changes_t - changes_p),
  510.                           float(i) / iterations * 100, end_l - start_l))
  511.         changes_p, changes_t = 0, 0
  512.         start_l = time.time()
  513.  
  514.     # calculate ED every 50000 iterations
  515.     if i % 50000 == 0 and i != 0:
  516.         ed = compare(buffer1=sourceBuffer, buffer2=tempBuffer, pixels=sourceBuffer)
  517.         if abs(pre_ed - ed) < ED_RoC and ED_RoC != 0.0:
  518.             break
  519.         print('Current ED {:.10f} ED RoC {:.10f}'.format(ed, abs(pre_ed - ed)))
  520.         pre_ed = ed
  521.  
  522.     # save around 50 images in progress to a dir for creating webm
  523.     if i % (iterations / 50) == 0 and i != 0:
  524.         BufferDraw(tempBuffer)
  525.         tempImage.save('{}/{:04d}.jpg'.format(dirname, counter))
  526.         ed = compare(buffer1=sourceBuffer, buffer2=tempBuffer, pixels=sourceBuffer)
  527.         print('Image {:04d} saved, ED {:.10f}'.format(counter, ed))
  528.         counter += 1
  529.  
  530. # save final result
  531. BufferDraw(tempBuffer)
  532. tempImage.save(outfile)
  533. ed = compare(buffer1=sourceBuffer, buffer2=tempBuffer, pixels=sourceBuffer)
  534. print('Image {} saved, ED {:.10f}.'.format(outfile, ed,))
  535.  
  536. end = time.time()
  537. timetaken = end - start
  538. print('Job done. {} iterations ({} accepted, {} discarded). It took {:.0f} seconds'.format(iterations, changes, iterations - changes, timetaken))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement