Software Público Brasileiro

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

"""
Author: Caio Marcelo Campoy Guedes
E-Mail: caiomcg@gmail.com

Author: Erickson Silva
E-Mail: erickson.silva@lavid.ufpb.br

Author: Jorismar Barbosa
E-Mail: jorismar.barbosa@lavid.ufpb.br

Author: Wesnydy Lima Ribeiro
E-Mail: wesnydy@lavid.ufpb.br
"""

import json
import logging
import os
import pika
import PikaManager
import subprocess

from thread import start_new_thread
from time import sleep
from urllib import urlretrieve

def make_dir_if_exists(path):
    if not os.path.exists(path):
        os.makedirs(path)

# Logging configuration.
logger = logging.getLogger('mixer')
logger.setLevel(logging.DEBUG)

fh = logging.FileHandler('/home/vlibras/log/mixer.log')
fh.setLevel(logging.DEBUG)

ch = logging.StreamHandler()
ch.setLevel(logging.INFO)

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
ch.setFormatter(formatter)

logger.addHandler(fh)
logger.addHandler(ch)

# Manager of queues connections.
#manager = PikaManager.PikaManager("150.165.205.10", "test", "test")

manager = PikaManager.PikaManager("rabbit")

PATH_MIXED_VIDEO = os.getenv("VLIBRAS_VIDEO_MIXED")
#Create Path if Needed
make_dir_if_exists(PATH_MIXED_VIDEO)

def main_video_height(main_video):
    """
    Extract height information of video.

    Parameters
    ----------
    main_video : string
        Video to obtain height.

    Returns
    -------
    string
        None if failed to extract info. The height if extraction has been successfuly.
    """
    logger.info("Extracting resolution of main video")
    try:
        # Obtains the main video height using ffprobe
        ffprobe = subprocess.Popen(
            [
                "ffprobe",
                 "-loglevel", "error",
                 "-select_streams", "v:0",
                 "-print_format", "json",
                 "-show_entries", 'stream=height',
                 main_video
            ],
            stdout=subprocess.PIPE,
            shell=False
        )
        # The results comes in a json
        video_height = json.loads(ffprobe.communicate()[0])
        # Returns the height obtained
        return video_height['streams'][0]['height']
    except:
        logger.error("Error when extracting resolution, default will be used")
        return None

def secondary_video_heigth(main_height, window_size):
    """
    Calculates the height of window.

    Parameters
    ----------
    main_height : string
        Height of main video.
    window_size : string
        User choice of size of window.

    Returns
    -------
    number
        The height of window.
    """
    logger.info("Calculating the resolution of the libras window")
    # Set the default height of main video if a height is not given
    if main_height is None:
        main_height = 324
    # Calculates the height of the small window
    if window_size == 'small':
        return int(0.3 * int(main_height))
    # Calculates the height of the large window
    elif window_size == 'large':
        return int(0.5 * int(main_height))
    # Calculates the height of the medium window (default)
    else:
        return int(0.4 * int(main_height))

def secondary_video_position(window_position):
    """
    Defines the position of window.

    Parameters
    ----------
    window_position : string
        User choice of position of window.

    Returns
    -------
    string
        The configurations of position of window.
    """
    logger.info("Defining the position of the libras window")
    # Overlap the window at top Left on main video
    if window_position == 'top_left':
        return "10:10"
    # Overlap the window at top right on main video
    elif window_position == 'top_right':
        return "main_w-overlay_w-10:10"
    # Overlap the window at bottom left on main video
    elif window_position == 'bottom_left':
        return "10:main_h-overlay_h-10"
    # Overlap the window at bottom right on main video (default)
    else:
        return "main_w-overlay_w-10:main_h-overlay_h-10"

def run(ch, method, properties, body):
    """
    Execute the worker.

    Parameters
    ----------
    ch : object
        Channel of communication.
    method : function
        Callback method.
    properties : object
        Message containing a set of 14 properties.
    body : string
        Json string containing the necessary arguments for workers.
    """
    logger.info("Processing request " + properties.correlation_id.encode("utf-8"))
    body = json.loads(body)
    try:
        logger.info("Downloading main video")
        main_video = urlretrieve(body["video"].encode("utf-8"))[0]
    except IOError as ex:
        logger.error("Download of video fail")
        return
    # Get the main video height
    main_height = main_video_height(main_video)
    # Calculates the window height based on the main video height
    window_heigth = secondary_video_heigth(main_height, body["window_size"].encode("utf-8"))
    # The width is proportionally to height (represented by -1)
    window_width = '-1'
    # Get the window position regarding to the main video
    window_pos = secondary_video_position(body["window_position"].encode("utf-8"))
    # Defines the window movie
    movie = 'movie=' + body["libras-video"].encode("utf-8")
    # Defines the scale of window movie
    scale = 'scale=' + str(window_width) + ':' + str(window_heigth)
    # Defines the overlay position
    overlay = '[movie] overlay=' + window_pos + ' [out]'
    # -Vf param
    filter_graph = ','.join([movie, scale, 'setpts=PTS-STARTPTS', overlay])
    # Generates the output file path
    mixed_video = os.path.join(PATH_MIXED_VIDEO, properties.correlation_id.encode("utf-8")+".mp4")
    # Mix videos using ffmpeg
    print ("Mixing videos...")
    logger.info("Mixing videos")
    try:
        subprocess.call(
            [
                "ffmpeg",
                "-loglevel", "error",
                "-i", main_video,
                "-y",
                "-vf", filter_graph,
                "-qscale", "0",
                "-strict", "experimental",
                "-vcodec", "libx264",
                "-preset", "fast",
                "-r", "30",
                "-threads", "4",
                mixed_video
            ],
            shell=False
        )
        logger.info("Mixing successfuly")
    except OSError as ex:
        logger.error("Mixing fail")
        print ("Error")
    # Add mixed video to the body
    body["mixed_video"] = mixed_video

    logger.info("Cleaning temp files")
    os.remove(main_video)
    os.remove(body["libras-video"])

    logger.info("Sending mixed video to the videos queue")
    manager.send_to_queue("videos", body, properties)
    print ("Ok")

def keep_alive(conn_send, conn_receive):
	while True:
		sleep(30)
		try:
			conn_send.process_data_events()
			conn_receive.process_data_events()
		except:
			continue

start_new_thread(keep_alive, (manager.get_conn_send(), manager.get_conn_receive()))

# Starts listening
print("Mixer listening...")
while True:
    try:
        manager.receive_from_queue("libras", run)
    except KeyboardInterrupt:
        manager.close_connections()
        os._exit(0)