Movie frame color rainbow (python 3)

movie_rainbow.py

"""
Loads a movie, analyzes the frames in a number of worker threads, and uses the average frame colors to create a rainbow image.
"""

import cv2
import threading
import queue
from os.path import basename
from multiprocessing import cpu_count
from PIL import Image
from numpy import array, uint8


def read_frames(path):
	cap = cv2.VideoCapture(path)
	framenr = 0
	while cap.isOpened():
		success, frame = cap.read()
		if not success:
			break
		yield frame
		framenr += 1
	print('read {} frames'.format(framenr))


class Worker(threading.Thread):
	def __init__(self, work, results, *args, **kwargs):
		self.work = work
		self.results = results
		super().__init__(*args, **kwargs)

	def run(self):
		try:
			while True:
				try:
					work = self.work.get(timeout=30)  # 3s timeout
				except queue.Empty:
					return
				if work is None:
					return
				self.results.put(self.process_frame(work))
				self.work.task_done()
		finally:
			self.work.task_done()

	def process_frame(self, frame):
		return frame.sum((0, 1)) / (frame.shape[0] * frame.shape[1])


# Movie file
pth = 'the-duck-song.mp4'

# Start workers
worker_count = max(cpu_count() - 1, 1)
print('{} workers'.format(worker_count))
work_queue = queue.Queue(maxsize=worker_count * 2)
result_queue = queue.Queue(100_000_000)
workers = []
for workernr in range(worker_count):
	worker = Worker(work_queue, result_queue)
	worker.start()
	workers.append(worker)

# Add all the frames
for nr, frame in enumerate(read_frames(pth)):
	if nr % 50 == 0:
		print('read frame {:d}'.format(nr))
	work_queue.put(frame, timeout=60)

# Collect the results
averages = []
try:
	while True:
		averages.append(result_queue.get_nowait())
		result_queue.task_done()
except queue.Empty:
	pass
result_queue.join()

# Stop all the workers
for worker in workers:
	work_queue.put(None)
work_queue.join()
for worker in workers:
	worker.join()
print('workers stopped')

print('{} averages computed'.format(len(averages)))

# Show and store the image
print('preparing image')
data = array(list(averages for h in range(100)))
print('data shape: {} max: {}'.format(data.shape, data[0, :, :].mean(-1).max()))
print(data)
img = Image.fromarray(data.astype(uint8), 'RGB')
img.show()
img.save(basename(pth) + '.png', 'PNG')

print('done')

Curious whether this only copies the reference (shares the memory of the movie frames), or does pickling like multiprocessing.Pool().map() does...