4
\$\begingroup\$

Here is a snippet of code from a project I have been working on that can be found here. I am looking for suggestions on how I can improve my code to make it cleaner and maybe improve performance.

from PIL import Image, ImageDraw
import subprocess as sp
import numpy
from pathlib import Path
import argparse
import AverageColor

def usage():
    print("USAGE: videoparser.py --file {FILENAME} -f --FRAME {NUMBER OF FRAMES TO SKIP")


def get_resolution(path):
    out = sp.check_output(["ffprobe", "-v", "error", "-of", "flat=s=_",
        "-select_streams", "v:0", "-show_entries", "stream=height,width", path])

    lines = out.split("\n")
    vidwidth = lines[0].split('=')[1]
    vidheight = lines[1].split('=')[1]
    return vidwidth, vidheight

# Get the arguments from the command line and assign them to variables
parser = argparse.ArgumentParser(description="Analyze the change in colors of videos over time")
parser.add_argument('-f', "--frames", dest="FRAMESKIPCOUNT", default=24)
parser.add_argument('--file', dest="FILENAME", required=True)
args = vars(parser.parse_args())

FRAME_SKIP_COUNT = args['FRAMESKIPCOUNT']
FILENAME = args['FILENAME']
if args['FRAMESKIPCOUNT']:
    FRAME_SKIP_COUNT = int(args['FRAMESKIPCOUNT'])

# Make sure the file is valid
my_file = Path(FILENAME)
if not my_file.is_file():
    raise FileNotFoundError('File was not found', args['FILENAME'])

# Run the command to start the FFMPEG library
FFMPEG_BIN = "ffmpeg.exe"

command = [FFMPEG_BIN,
           '-i', FILENAME,
           '-f', 'image2pipe',
           '-pix_fmt', 'rgb24',
   '-vcodec', 'rawvideo', '-']

pipe = sp.Popen(command, stdout=sp.PIPE, bufsize=10 ** 8, shell=True)  # Open pipe to start receiving pixel data
i = 0  # i is for processing every 'FRAME_SKIP_COUNT'th element
pos = 0  # to hold the position to draw the next line
width, height = get_resolution(FILENAME)
# Declare a blank image and prepare it for drawing
finalImage = Image.new("RGB", (20000, height))
finalImageDraw = ImageDraw.Draw(finalImage, "RGB")

while True:
    # read 1 frame
    raw_image = pipe.stdout.read(1280*720*3)

    if i % FRAME_SKIP_COUNT == 0:
        # transform the byte read into a numpy array
        image = numpy.fromstring(raw_image, dtype='uint8')
        if image.size == 0:
            # No more data, reached the end of the video
            break
        # Put the data into an image for processing (Can this be skipped to improve performance?)
        image = image.reshape((720, 1280, 3))
        img = Image.fromarray(image, 'RGB')

        avgColor = AverageColor.averagecolorfromimage(img)
        # Draw a line
        finalImageDraw.line([(pos, 0), (pos, height)], fill="rgb" + str(avgColor))
        pos += 1
    # throw away the data in the pipe's buffer.
    pipe.stdout.flush()
    i += 1
# Crop the image and save it
finalImage = finalImage.crop([0, 0, pos, 1000])
finalImage.save(FILENAME + " - color.png")

Here is the AverageColor implementation

from pathlib import Path

from PIL import Image


def averagecolorfromfile(file):

    file = Path(file)
    if not file.is_file():
        raise FileNotFoundError('File was not found', file)

    im = Image.open(file)

    # create list of pixel's RGB values and their count (stored in color[0])
    colors = im.getcolors(1000000)

    im.close()

    count = 0
    c1 = 0
    c2 = 0
    c3 = 0

    # add all the values up and divide by the total number of
    # pixels to get the average rgb value

    for color in colors:
        count += color[0]
        c1 += (color[1][0] * color[0])
        c2 += (color[1][1] * color[0])
        c3 += (color[1][2] * color[0])

    c1 /= count
    c2 /= count
    c3 /= count

    return int(c1), int(c2), int(c3)


def averagecolorfromimage(image):

    # create list of pixel's RGB values and their count (stored in color[0])
    colors = image.getcolors(1000000)

    count = 0
    c1 = 0
    c2 = 0
    c3 = 0

    # add all the values up and divide by the total number of
    # pixels to get the average rgb value

    for color in colors:
        count += color[0]
        c1 += (color[1][0] * color[0])
        c2 += (color[1][1] * color[0])
        c3 += (color[1][2] * color[0])

    c1 /= count
    c2 /= count
    c3 /= count

    return int(c1), int(c2), int(c3)
\$\endgroup\$
2
\$\begingroup\$

Here's an alternative implementation of averagecolorfromfile:

import scipy.ndimage

def averagecolorfromfile2(file):
    "Return mean color of the pixels in the image loaded from file."
    return scipy.ndimage.imread(file).mean(axis=(0, 1)).round().astype('uint8')

I find that this is about 16 times as fast as the code in the post:

>>> data = np.random.randint(0, 256, size=(1000, 1000, 3), dtype='uint8')
>>> filename = 'cr153189.png'
>>> PIL.Image.fromarray(data).save(filename)
>>> from timeit import timeit
>>> timeit(lambda:averagecolorfromfile(filename), number=1)
1.599527531012427
>>> timeit(lambda:averagecolorfromfile2(filename), number=1)
0.09920550498645753

Note also that averagecolorfromfile returns a truncated result (that is, rounded towards zero). But for most purposes it would be more accurate to round to the nearest integer, as in averagecolorfromfile2:

>>> averagecolorfromfile(filename)
(127, 127, 127)
>>> averagecolorfromfile2(filename)
array([127, 128, 128], dtype=uint8)
\$\endgroup\$
  • \$\begingroup\$ Could you elaborate on what scipy.ndimage.imread(file) is doing, or point me to some material to read. I am trying to wrap my head around it but I don't quite understand why it works. \$\endgroup\$ – Jacob Malachowski Jan 24 '17 at 2:06
  • \$\begingroup\$ @JacobMalachowski: See the documentation. \$\endgroup\$ – Gareth Rees Jan 24 '17 at 12:15

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.