distortioner/cas.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from argparse import ArgumentParser
from tempfile import TemporaryDirectory
import logging
import cv2
import asyncio
from pathlib import Path
from wand.image import Image

log = logging.getLogger('distortioner')

ap = ArgumentParser(
    description = 'a script that makes memetic distorted videos'
)

class TicketedDict(dict):
    def __init__(self, *args, **kwargs):
        self._ni = 0
        self._log = logging.getLogger('TicketedDict')
        self.event = asyncio.Event()
        super().__init__(*args, **kwargs)

    def has_next(self):
        self._log.debug("has_next is %s for %i", self._ni in self, self._ni)
        return self._ni in self

    async def wait(self):
        self._log.debug("requested wait")
        while not self.has_next():
            await self.event.wait()
            self._log.debug("wait broke")
            self.event.clear()

    async def pop(self):
        self._log.debug("requested pop")
        await self.wait()
        ret = super().pop(self._ni)
        self._ni += 1
        return ret

    def notify(self, *args, **kwargs):
        self._log.debug("requested notify")
        return self.event.set(*args, **kwargs)

ap.add_argument('input')
ap.add_argument('output')
ap.add_argument(
    '--distort-percentage','--distort-pct',  '-d', 
    default = 60.0,
    type = float,
    help = 'Percentage of image distortion.'
)
ap.add_argument(
    '--distort-percentage-end','--distort-pct-end',  '-D', 
    default = None,
    type = float,
    help = 'If specified, distortion percentage will gradually change towards specified percentage.'
)
ap.add_argument(
    '--distort-end','--distort',  '-E', 
    default = None,
    type = float,
    help = 'If specified, distortion change will happen only up to specific video progress.'
)
ap.add_argument(
    '--vibrato-frequency','--vibrato-freq',  '-f',
    default = 5.0,
    type = float,
    help = 'Modulation frequency in Hertz. Range is 0.1 - 20000.0. Default value is 5.0 Hz.'
)
ap.add_argument(
    '--vibrato-modulation-depth','--vibrato-depth',  '-m',
    default = 0.5,
    type = float,
    help = 'Depth of modulation as a percentage. Range is 0.0 - 1.0. Default value is 0.5.'
)
ap.add_argument(
    '--debug',
    action = 'store_true',
    default = False,
    help = 'Print debugging messages.'
)
def process_image(source, destination, distort):
    log.debug("distorting: src:'%s', dst:'%s' ", source, destination)
    with Image(filename=source) as original:
        dst_width = int(original.width*(distort / 100.))
        dst_height = int(original.height*(distort / 100.))
        log.debug("dst:'%s' size:%ix%i", destination, dst_width, dst_height)
        with original.clone() as distorted, \
             open(destination, mode='wb') as out:
            distorted.liquid_rescale(dst_width, dst_height)
            distorted.resize(original.width, original.height)
            distorted.save(out)

async def process_frames(coro, queue_in, out_pile):
    while True:
        frame_data = await queue_in.get()
        nr = frame_data.pop('nr')
        log.debug("processing frame '%s'", frame_data)
        await asyncio.to_thread(coro, **frame_data)
        queue_in.task_done()
        out_pile[nr]=frame_data['destination']
        log.debug("put item to output pile, notifying")
        out_pile.notify()
        log.debug("notified")
        

async def read_frames(capture, frames_distorted, frames_original, queue, tasks, distort_start, distort_end=None, distort_end_frame=None):
    frames_read = 0
    if distort_end_frame is None:
        distort_end_frame = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    while True:
        read_ok, frame = capture.read()
        if not read_ok:
            break
        log.debug("reading frame %i", frames_read)
        if distort_end is not None:
            distort = distort_start \
                + (distort_end - distort_start) \
                * (min(frames_read, distort_end_frame) / distort_end_frame ) 
        frame_filename = f'frame_{str(frames_read).zfill(32)}.png'
        frame_original = str(Path(frames_original)/frame_filename)
        frame_distorted = str(Path(frames_distorted)/frame_filename)
        cv2.imwrite(frame_original, frame)
        log.debug("saving frame %i: filename: %s", frames_read, frame_filename)
        await queue.put({
            'source': frame_original,
            'destination': frame_distorted,
            'distort': distort,
            'nr': frames_read
        })
        frames_read += 1
    log.info("waiting for queue to empty")
    await queue.join()
    log.info("quitting")
    for worker in tasks:
        worker.cancel()

async def write_frames(output, pile):
    while True:
        log.debug("getting next item ...")
        frame_distorted = await pile.pop()
        log.info("writing frame '%s'", frame_distorted)
        newframe = cv2.imread(frame_distorted)
        output.write(newframe)
        log.debug("finished frame '%s'", frame_distorted)


async def distort_video(capture, output, distort_start, distort_end=None, distort_end_frame=None):
    distort = distort_start
    video_width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    video_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    output_pile = TicketedDict()
    with \
            TemporaryDirectory() as frames_distorted, \
            TemporaryDirectory() as frames_original:
        log.debug(frames_distorted)
        log.debug(frames_original)
        log.debug("creating queues")
        capture_queue = asyncio.Queue(20)
        output_queue = asyncio.PriorityQueue()
        workers = [ asyncio.create_task(process_frames(process_image, capture_queue, output_pile)) for i in range(10) ]
        workers += [ asyncio.create_task(write_frames(output, output_pile)) ]
        generator = asyncio.create_task(read_frames(
            capture, frames_distorted, frames_original, capture_queue, workers, distort_start, distort_end, distort_end_frame
        ))
        await asyncio.gather(generator, *workers, return_exceptions=True)

    log.debug('done with distorting video frames')

def main():
    args = ap.parse_args()
    if args.debug:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
    log.debug(args)
    if args.distort_percentage <= 0.0:
        raise ValueError("--distort_percentage must be positive number")

    capture = cv2.VideoCapture(args.input)
    fps = capture.get(cv2.CAP_PROP_FPS)
    video_width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    video_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frames = args.distort_end or int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    log.info("video: name:'%s', fps:%i, frames:%i, size:%ix%i", args.input, fps, frames, video_width, video_height)
    output = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*'mp4v'), fps, (video_width, video_height))
    asyncio.run(distort_video(capture, output, args.distort_percentage, args.distort_percentage_end, frames-1))
    capture.release()
    output.release()


if __name__ == '__main__':
    main()