simencode.py

#!/usr/bin/env python3
#
# simtv - an Arduino TV Simulator
#
# Copyright (C) 2018 Richard "Shred" Körber
#   https://github.com/shred/simtv
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import argparse
import colorsys
import cv2
import math
import numpy


def streamVideo(cap, targetFps, maxFrames):
    """Streams the capture as one pixel (r,g,b) tuple per frame."""
    frameCount = 0
    targetFrameCount = 0
    while maxFrames is None or frameCount < maxFrames:
        ret, frame = cap.read()
        if not ret:
            break

        frameCount = cap.get(cv2.CAP_PROP_POS_FRAMES)
        position = cap.get(cv2.CAP_PROP_POS_MSEC)

        print('Frame: %7d %3d:%02d -> %7d\r' % (
            frameCount,
            (int) (position / 1000 / 60),
            (int) (position / 1000 % 60),
            targetFrameCount), end='', flush=True)

        if position >= targetFrameCount * 1000 / targetFps:
            pixel = numpy.average(numpy.average(frame, axis=0), axis=0)
            targetFrameCount += 1
            yield (pixel[2] / 256, pixel[1] / 256, pixel[0] / 256)

    print()


def streamPixels(file, targetFps, maxFrames):
    """Streams a video file as pixel stream with the given target FPS."""
    cap = cv2.VideoCapture(file)
    try:
        if cap.isOpened():
            for color in streamVideo(cap, targetFps, maxFrames):
                yield color
    finally:
        cap.release()


def streamColorData(stream, harddiff, softdiff, brightness):
    """Converts a pixel stream to a data stream that is readable for the Arduino sketch."""
    previousColor = None
    writtenColor = None
    count = 0

    for color in stream:
        if writtenColor is None:
            yield convertColor(0, color, brightness)
            writtenColor = color

        elif isColorCut(color, previousColor, harddiff / 3, harddiff):
            if count > 0:
                yield convertColor(count - 1, color, brightness)
            yield convertColor(0, color, brightness)
            writtenColor = color
            count = 0

        elif count > 2 and isColorCut(color, writtenColor, softdiff / 2, softdiff):
            yield convertColor(count, color, brightness)
            writtenColor = color
            count = 0

        elif count == 15:
            yield convertColor(count, color, brightness)
            writtenColor = color
            count = 0

        else:
            count += 1

        previousColor = color

    yield convertColor(count, previousColor, brightness)


def convertColor(step, color, brightness):
    """Returns a tuple of the color and step."""
    gcolor = [max(min(int(round(c * brightness * 16)), 15), 0) for c in color]
    return (step if step <= 15 else 15, gcolor[0], gcolor[1], gcolor[2])


def isColorCut(c1, c2, hueLimit, brightLimit):
    """Checks if color c1 and color c2 difference exceed the given limits."""
    hsv1 = colorsys.rgb_to_hsv(c1[0], c1[1], c1[2])
    hsv2 = colorsys.rgb_to_hsv(c2[0], c2[1], c2[2])

    # hue wraps around between 0.0 and 1.0
    hue1 = hsv1[0] + (1.0 if hsv1[0] < 0.5 and hsv2[0] >= 0.5 else 0.0)
    hue2 = hsv2[0] + (1.0 if hsv2[0] < 0.5 and hsv1[0] >= 0.5 else 0.0)
    dh = math.fabs(hue1 - hue2)

    db = math.fabs(hsv1[2] - hsv2[2])

    return dh >= hueLimit or db >= brightLimit


def writeSource(out, filename, fps, gamma, stream):
    """Writes out the stream.h source code."""
    out.write("""/*
 * Converted movie stream: %s
 *
 * Generated by simencode.py, do not edit!
 */

#include <avr/pgmspace.h>

const unsigned int FPS = %d;

PROGMEM const byte GAMMA[] = {
  """ % (filename, fps))

    for cv in range(16):
        if cv > 0:
            out.write(", ")
        out.write("%d" % (max(min(round(math.pow(cv / 15, gamma) * 255), 255), 0)))

    out.write("\n};\n\nPROGMEM const unsigned int DATA[] = {");

    count = 0
    for c in stream:
        if count > 0:
            out.write(", ")
        if count % 8 == 0:
            out.write("\n  ")
        out.write("0x%1X%1X%1X%1X" % c)
        count += 1

    out.write("\n};\n")



parser = argparse.ArgumentParser(description='Convert movie to SimTV source')
parser.add_argument('file',
            nargs='?',
            help='input movie file')
parser.add_argument('-o', '--out',
            dest='out',
            default='./simtv/stream.h',
            help='output source file (default: simtv/stream.h)')
parser.add_argument('-x', '--frames',
            dest='frames',
            type=int,
            help='maximum number of frames to convert')
parser.add_argument('-r', '--fps',
            type=int,
            default=10,
            choices=range(1, 51),
            metavar="[1-50]",
            dest='fps',
            help='target frames per second (default: 10)')
parser.add_argument('-g', '--gamma',
            dest='gamma',
            type=float,
            default=2.2,
            help='gamma correction (default: 2.2)')
parser.add_argument('-b', '--brightness',
            dest='brightness',
            type=float,
            default=1.0,
            help='brightness correction (default: 1.0)')
parser.add_argument('-S', '--soft',
            dest='soft',
            type=int,
            default=7,
            choices=range(0, 101),
            metavar="[0-100]",
            help='hue/brightness threshold for soft color cuts (default: 7%%)')
parser.add_argument('-H', '--hard',
            dest='hard',
            type=int,
            default=20,
            choices=range(0, 101),
            metavar="[0-100]",
            help='hue/brightness threshold for hard color cuts (default: 20%%)')
args = parser.parse_args()

with open(args.out, 'w') as w:
     writeSource(w, args.file, args.fps, args.gamma, \
        streamColorData( \
            streamPixels(args.file, args.fps, args.frames), \
            args.hard / 100, \
            args.soft / 100, \
            args.brightness))