File: afy/arguments.py

from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument("--config", help="path to config")
parser.add_argument("--checkpoint", default='vox-cpk.pth.tar', help="path to checkpoint to restore")

parser.add_argument("--relative", dest="relative", action="store_true", help="use relative or absolute keypoint coordinates")
parser.add_argument("--adapt_scale", dest="adapt_scale", action="store_true", help="adapt movement scale based on convex hull of keypoints")
parser.add_argument("--no-pad", dest="no_pad", action="store_true", help="don't pad output image")
parser.add_argument("--enc_downscale", default=1, type=float, help="Downscale factor for encoder input. Improves performance with cost of quality.")

parser.add_argument("--virt-cam", type=int, default=0, help="Virtualcam device ID")
parser.add_argument("--no-stream", action="store_true", help="On Linux, force no streaming")

parser.add_argument("--verbose", action="store_true", help="Print additional information")
parser.add_argument("--hide-rect", action="store_true", default=False, help="Hide the helper rectangle in preview window")

parser.add_argument("--avatars", default="./avatars", help="path to avatars directory")

parser.add_argument("--is-worker", action="store_true", help="Whether to run this process as a remote GPU worker")
parser.add_argument("--is-client", action="store_true", help="Whether to run this process as a client")
parser.add_argument("--in-port", type=int, default=5557, help="Remote worker input port")
parser.add_argument("--out-port", type=int, default=5558, help="Remote worker output port")
parser.add_argument("--in-addr", type=str, default=None, help="Socket address for incoming messages, like example.com:5557")
parser.add_argument("--out-addr", type=str, default=None, help="Socker address for outcoming messages, like example.com:5558")
parser.add_argument("--jpg_quality", type=int, default=95, help="Jpeg copression quality for image transmission")

parser.set_defaults(relative=False)
parser.set_defaults(adapt_scale=False)
parser.set_defaults(no_pad=False)

opt = parser.parse_args()

if opt.is_client and (opt.in_addr is None or opt.out_addr is None):
    raise ValueError("You have to set --in-addr and --out-addr")



File: afy/cam_fomm.py

import os, sys
from sys import platform as _platform
import glob
import yaml
import time
import requests

import numpy as np
import cv2

from afy.videocaptureasync import VideoCaptureAsync
from afy.arguments import opt
from afy.utils import info, Once, Tee, crop, pad_img, resize, TicToc
import afy.camera_selector as cam_selector


log = Tee('./var/log/cam_fomm.log')

# Where to split an array from face_alignment to separate each landmark
LANDMARK_SLICE_ARRAY = np.array([17, 22, 27, 31, 36, 42, 48, 60])

if _platform == 'darwin':
    if not opt.is_client:
        info('\nOnly remote GPU mode is supported for Mac (use --is-client and --connect options to connect to the server)')
        info('Standalone version will be available lately!\n')
        exit()


def is_new_frame_better(source, driving, predictor):
    global avatar_kp
    global display_string
    
    if avatar_kp is None:
        display_string = "No face detected in avatar."
        return False
    
    if predictor.get_start_frame() is None:
        display_string = "No frame to compare to."
        return True
    
    driving_smaller = resize(driving, (128, 128))[..., :3]
    new_kp = predictor.get_frame_kp(driving)
    
    if new_kp is not None:
        new_norm = (np.abs(avatar_kp - new_kp) ** 2).sum()
        old_norm = (np.abs(avatar_kp - predictor.get_start_frame_kp()) ** 2).sum()
        
        out_string = "{0} : {1}".format(int(new_norm * 100), int(old_norm * 100))
        display_string = out_string
        log(out_string)
        
        return new_norm < old_norm
    else:
        display_string = "No face found!"
        return False


def load_stylegan_avatar():
    url = "https://thispersondoesnotexist.com/image"
    r = requests.get(url, headers={'User-Agent': "My User Agent 1.0"}).content

    image = np.frombuffer(r, np.uint8)
    image = cv2.imdecode(image, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    image = resize(image, (IMG_SIZE, IMG_SIZE))

    return image

def load_images(IMG_SIZE = 256):
    avatars = []
    filenames = []
    images_list = sorted(glob.glob(f'{opt.avatars}/*'))
    for i, f in enumerate(images_list):
        if f.endswith('.jpg') or f.endswith('.jpeg') or f.endswith('.png'):
            img = cv2.imread(f)
            if img is None:
                log("Failed to open image: {}".format(f))
                continue

            if img.ndim == 2:
                img = np.tile(img[..., None], [1, 1, 3])
            img = img[..., :3][..., ::-1]
            img = resize(img, (IMG_SIZE, IMG_SIZE))
            avatars.append(img)
            filenames.append(f)
    return avatars, filenames

def change_avatar(predictor, new_avatar):
    global avatar, avatar_kp, kp_source
    avatar_kp = predictor.get_frame_kp(new_avatar)
    kp_source = None
    avatar = new_avatar
    predictor.set_source_image(avatar)


def draw_rect(img, rw=0.6, rh=0.8, color=(255, 0, 0), thickness=2):
    h, w = img.shape[:2]
    l = w * (1 - rw) // 2
    r = w - l
    u = h * (1 - rh) // 2
    d = h - u
    img = cv2.rectangle(img, (int(l), int(u)), (int(r), int(d)), color, thickness)

def kp_to_pixels(arr):
    '''Convert normalized landmark locations to screen pixels'''
    return ((arr + 1) * 127).astype(np.int32)

def draw_face_landmarks(img, face_kp, color=(20, 80, 255)):
    if face_kp is not None:
        img = cv2.polylines(img, np.split(kp_to_pixels(face_kp), LANDMARK_SLICE_ARRAY), False, color)

def print_help():
    info('\n\n=== Control keys ===')
    info('1-9: Change avatar')
    for i, fname in enumerate(avatar_names):
        key = i + 1
        name = fname.split('/')[-1]
        info(f'{key}: {name}')
    info('W: Zoom camera in')
    info('S: Zoom camera out')
    info('A: Previous avatar in folder')
    info('D: Next avatar in folder')
    info('Q: Get random avatar')
    info('X: Calibrate face pose')
    info('I: Show FPS')
    info('ESC: Quit')
    info('\nFull key list: https://github.com/alievk/avatarify#controls')
    info('\n\n')


def draw_fps(frame, fps, timing, x0=10, y0=20, ystep=30, fontsz=0.5, color=(255, 255, 255)):
    frame = frame.copy()
    cv2.putText(frame, f"FPS: {fps:.1f}", (x0, y0 + ystep * 0), 0, fontsz * IMG_SIZE / 256, color, 1)
    cv2.putText(frame, f"Model time (ms): {timing['predict']:.1f}", (x0, y0 + ystep * 1), 0, fontsz * IMG_SIZE / 256, color, 1)
    cv2.putText(frame, f"Preproc time (ms): {timing['preproc']:.1f}", (x0, y0 + ystep * 2), 0, fontsz * IMG_SIZE / 256, color, 1)
    cv2.putText(frame, f"Postproc time (ms): {timing['postproc']:.1f}", (x0, y0 + ystep * 3), 0, fontsz * IMG_SIZE / 256, color, 1)
    return frame


def draw_landmark_text(frame, thk=2, fontsz=0.5, color=(0, 0, 255)):
    frame = frame.copy()
    cv2.putText(frame, "ALIGN FACES", (60, 20), 0, fontsz * IMG_SIZE / 255, color, thk)
    cv2.putText(frame, "THEN PRESS X", (60, 245), 0, fontsz * IMG_SIZE / 255, color, thk)
    return frame


def draw_calib_text(frame, thk=2, fontsz=0.5, color=(0, 0, 255)):
    frame = frame.copy()
    cv2.putText(frame, "FIT FACE IN RECTANGLE", (40, 20), 0, fontsz * IMG_SIZE / 255, color, thk)
    cv2.putText(frame, "W - ZOOM IN", (60, 40), 0, fontsz * IMG_SIZE / 255, color, thk)
    cv2.putText(frame, "S - ZOOM OUT", (60, 60), 0, fontsz * IMG_SIZE / 255, color, thk)
    cv2.putText(frame, "THEN PRESS X", (60, 245), 0, fontsz * IMG_SIZE / 255, color, thk)
    return frame


def select_camera(config):
    cam_config = config['cam_config']
    cam_id = None

    if os.path.isfile(cam_config):
        with open(cam_config, 'r') as f:
            cam_config = yaml.load(f, Loader=yaml.FullLoader)
            cam_id = cam_config['cam_id']
    else:
        cam_frames = cam_selector.query_cameras(config['query_n_cams'])

        if cam_frames:
            if len(cam_frames) == 1:
                cam_id = list(cam_frames)[0]
            else:
                cam_id = cam_selector.select_camera(cam_frames, window="CLICK ON YOUR CAMERA")
            log(f"Selected camera {cam_id}")

            with open(cam_config, 'w') as f:
                yaml.dump({'cam_id': cam_id}, f)
        else:
            log("No cameras are available")

    return cam_id


if __name__ == "__main__":
    with open('config.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    global display_string
    display_string = ""

    IMG_SIZE = 256

    log('Loading Predictor')
    predictor_args = {
        'config_path': opt.config,
        'checkpoint_path': opt.checkpoint,
        'relative': opt.relative,
        'adapt_movement_scale': opt.adapt_scale,
        'enc_downscale': opt.enc_downscale
    }
    if opt.is_worker:
        from afy import predictor_worker
        predictor_worker.run_worker(opt.in_port, opt.out_port)
        sys.exit(0)
    elif opt.is_client:
        from afy import predictor_remote
        try:
            predictor = predictor_remote.PredictorRemote(
                in_addr=opt.in_addr, out_addr=opt.out_addr,
                **predictor_args
            )
        except ConnectionError as err:
            log(err)
            sys.exit(1)
        predictor.start()
    else:
        from afy import predictor_local
        predictor = predictor_local.PredictorLocal(
            **predictor_args
        )

    cam_id = select_camera(config)

    if cam_id is None:
        exit(1)

    cap = VideoCaptureAsync(cam_id)
    cap.start()

    avatars, avatar_names = load_images()

    enable_vcam = not opt.no_stream

    ret, frame = cap.read()
    stream_img_size = frame.shape[1], frame.shape[0]

    if enable_vcam:
        if _platform in ['linux', 'linux2']:
            try:
                import pyfakewebcam
            except ImportError:
                log("pyfakewebcam is not installed.")
                exit(1)

            stream = pyfakewebcam.FakeWebcam(f'/dev/video{opt.virt_cam}', *stream_img_size)
        else:
            enable_vcam = False
            # log("Virtual camera is supported only on Linux.")
        
        # if not enable_vcam:
            # log("Virtual camera streaming will be disabled.")

    cur_ava = 0    
    avatar = None
    change_avatar(predictor, avatars[cur_ava])
    passthrough = False

    cv2.namedWindow('cam', cv2.WINDOW_GUI_NORMAL)
    cv2.moveWindow('cam', 500, 250)

    frame_proportion = 0.9
    frame_offset_x = 0
    frame_offset_y = 0

    overlay_alpha = 0.0
    preview_flip = False
    output_flip = False
    find_keyframe = False
    is_calibrated = False

    show_landmarks = False

    fps_hist = []
    fps = 0
    show_fps = False

    print_help()

    try:
        while True:
            tt = TicToc()

            timing = {
                'preproc': 0,
                'predict': 0,
                'postproc': 0
            }

            green_overlay = False

            tt.tic()

            ret, frame = cap.read()
            if not ret:
                log("Can't receive frame (stream end?). Exiting ...")
                break

            frame = frame[..., ::-1]
            frame_orig = frame.copy()

            frame, (frame_offset_x, frame_offset_y) = crop(frame, p=frame_proportion, offset_x=frame_offset_x, offset_y=frame_offset_y)

            frame = resize(frame, (IMG_SIZE, IMG_SIZE))[..., :3]

            if find_keyframe:
                if is_new_frame_better(avatar, frame, predictor):
                    log("Taking new frame!")
                    green_overlay = True
                    predictor.reset_frames()

            timing['preproc'] = tt.toc()

            if passthrough:
                out = frame
            elif is_calibrated:
                tt.tic()
                out = predictor.predict(frame)
                if out is None:
                    log('predict returned None')
                timing['predict'] = tt.toc()
            else:
                out = None

            tt.tic()
            
            key = cv2.waitKey(1)

            if cv2.getWindowProperty('cam', cv2.WND_PROP_VISIBLE) < 1.0:
                break
            elif is_calibrated and cv2.getWindowProperty('avatarify', cv2.WND_PROP_VISIBLE) < 1.0:
                break

            if key == 27: # ESC
                break
            elif key == ord('d'):
                cur_ava += 1
                if cur_ava >= len(avatars):
                    cur_ava = 0
                passthrough = False
                change_avatar(predictor, avatars[cur_ava])
            elif key == ord('a'):
                cur_ava -= 1
                if cur_ava < 0:
                    cur_ava = len(avatars) - 1
                passthrough = False
                change_avatar(predictor, avatars[cur_ava])
            elif key == ord('w'):
                frame_proportion -= 0.05
                frame_proportion = max(frame_proportion, 0.1)
            elif key == ord('s'):
                frame_proportion += 0.05
                frame_proportion = min(frame_proportion, 1.0)
            elif key == ord('H'):
                frame_offset_x -= 1
            elif key == ord('h'):
                frame_offset_x -= 5
            elif key == ord('K'):
                frame_offset_x += 1
            elif key == ord('k'):
                frame_offset_x += 5
            elif key == ord('J'):
                frame_offset_y -= 1
            elif key == ord('j'):
                frame_offset_y -= 5
            elif key == ord('U'):
                frame_offset_y += 1
            elif key == ord('u'):
                frame_offset_y += 5
            elif key == ord('Z'):
                frame_offset_x = 0
                frame_offset_y = 0
                frame_proportion = 0.9
            elif key == ord('x'):
                predictor.reset_frames()

                if not is_calibrated:
                    cv2.namedWindow('avatarify', cv2.WINDOW_GUI_NORMAL)
                    cv2.moveWindow('avatarify', 600, 250)
                
                is_calibrated = True
                show_landmarks = False
            elif key == ord('z'):
                overlay_alpha = max(overlay_alpha - 0.1, 0.0)
            elif key == ord('c'):
                overlay_alpha = min(overlay_alpha + 0.1, 1.0)
            elif key == ord('r'):
                preview_flip = not preview_flip
            elif key == ord('t'):
                output_flip = not output_flip
            elif key == ord('f'):
                find_keyframe = not find_keyframe
            elif key == ord('o'):
                show_landmarks = not show_landmarks
            elif key == ord('q'):
                try:
                    log('Loading StyleGAN avatar...')
                    avatar = load_stylegan_avatar()
                    passthrough = False
                    change_avatar(predictor, avatar)
                except:
                    log('Failed to load StyleGAN avatar')
            elif key == ord('l'):
                try:
                    log('Reloading avatars...')
                    avatars, avatar_names = load_images()
                    passthrough = False
                    log("Images reloaded")
                except:
                    log('Image reload failed')
            elif key == ord('i'):
                show_fps = not show_fps
            elif 48 < key < 58:
                cur_ava = min(key - 49, len(avatars) - 1)
                passthrough = False
                change_avatar(predictor, avatars[cur_ava])
            elif key == 48:
                passthrough = not passthrough
            elif key != -1:
                log(key)

            if overlay_alpha > 0:
                preview_frame = cv2.addWeighted( avatar, overlay_alpha, frame, 1.0 - overlay_alpha, 0.0)
            else:
                preview_frame = frame.copy()

            if show_landmarks:
                # Dim the background to make it easier to see the landmarks
                preview_frame = cv2.convertScaleAbs(preview_frame, alpha=0.5, beta=0.0)

                draw_face_landmarks(preview_frame, avatar_kp, (200, 20, 10))
                frame_kp = predictor.get_frame_kp(frame)
                draw_face_landmarks(preview_frame, frame_kp)
            
            if preview_flip:
                preview_frame = cv2.flip(preview_frame, 1)
                
            if green_overlay:
                green_alpha = 0.8
                overlay = preview_frame.copy()
                overlay[:] = (0, 255, 0)
                preview_frame = cv2.addWeighted( preview_frame, green_alpha, overlay, 1.0 - green_alpha, 0.0)

            timing['postproc'] = tt.toc()
                
            if find_keyframe:
                preview_frame = cv2.putText(preview_frame, display_string, (10, 220), 0, 0.5 * IMG_SIZE / 256, (255, 255, 255), 1)

            if show_fps:
                preview_frame = draw_fps(preview_frame, fps, timing)

            if not is_calibrated:
                preview_frame = draw_calib_text(preview_frame)
            elif show_landmarks:
                preview_frame = draw_landmark_text(preview_frame)

            if not opt.hide_rect:
                draw_rect(preview_frame)

            cv2.imshow('cam', preview_frame[..., ::-1])

            if out is not None:
                if not opt.no_pad:
                    out = pad_img(out, stream_img_size)

                if output_flip:
                    out = cv2.flip(out, 1)

                if enable_vcam:
                    out = resize(out, stream_img_size)
                    stream.schedule_frame(out)

                cv2.imshow('avatarify', out[..., ::-1])

            fps_hist.append(tt.toc(total=True))
            if len(fps_hist) == 10:
                fps = 10 / (sum(fps_hist) / 1000)
                fps_hist = []
    except KeyboardInterrupt:
        log("main: user interrupt")

    log("stopping camera")
    cap.stop()

    cv2.destroyAllWindows()

    if opt.is_client:
        log("stopping remote predictor")
        predictor.stop()

    log("main: exit")



File: afy/videocaptureasync.py

# https://github.com/gilbertfrancois/video-capture-async

import threading
import cv2
import time


WARMUP_TIMEOUT = 10.0


class VideoCaptureAsync:
    def __init__(self, src=0, width=640, height=480):
        self.src = src

        self.cap = cv2.VideoCapture(self.src)
        if not self.cap.isOpened():
            raise RuntimeError("Cannot open camera")

        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
        self.grabbed, self.frame = self.cap.read()
        self.started = False
        self.read_lock = threading.Lock()

    def set(self, var1, var2):
        self.cap.set(var1, var2)

    def isOpened(self):
        return self.cap.isOpened()

    def start(self):
        if self.started:
            print('[!] Asynchronous video capturing has already been started.')
            return None
        self.started = True
        self.thread = threading.Thread(target=self.update, args=(), daemon=True)
        self.thread.start()

        # (warmup) wait for the first successfully grabbed frame
        warmup_start_time = time.time()
        while not self.grabbed:
            warmup_elapsed_time = (time.time() - warmup_start_time)
            if warmup_elapsed_time > WARMUP_TIMEOUT:
                raise RuntimeError(f"Failed to succesfully grab frame from the camera (timeout={WARMUP_TIMEOUT}s). Try to restart.")

            time.sleep(0.5)

        return self

    def update(self):
        while self.started:
            grabbed, frame = self.cap.read()
            if not grabbed or frame is None or frame.size == 0:
                continue
            with self.read_lock:
                self.grabbed = grabbed
                self.frame = frame

    def read(self):
        while True:
            with self.read_lock:
                frame = self.frame.copy()
                grabbed = self.grabbed
            break
        return grabbed, frame

    def stop(self):
        self.started = False
        self.thread.join()

    def __exit__(self, exec_type, exc_value, traceback):
        self.cap.release()



File: afy/networking.py

import zmq
import numpy as np
import msgpack
import msgpack_numpy as m
m.patch()

from afy.utils import log


def check_connection(socket, timeout=1000):
    old_rcvtimeo = socket.RCVTIMEO
    socket.RCVTIMEO = timeout

    try:
        data = msgpack.packb(([], {}))
        socket.send_data('hello', data)
        attr_recv, data_recv = socket.recv_data()
        response = msgpack.unpackb(data_recv)
    except zmq.error.Again:
        return False
    finally:
        socket.RCVTIMEO = old_rcvtimeo

    log(f"Response to hello is {response}")
  
    return response == 'OK'


class SerializingSocket(zmq.Socket):
    """Numpy array serialization methods.

    Based on https://github.com/jeffbass/imagezmq/blob/master/imagezmq/imagezmq.py#L291

    Used for sending / receiving OpenCV images, which are Numpy arrays.
    Also used for sending / receiving jpg compressed OpenCV images.
    """

    def send_array(self, A, msg='NoName', flags=0, copy=True, track=False):
        """Sends a numpy array with metadata and text message.

        Sends a numpy array with the metadata necessary for reconstructing
        the array (dtype,shape). Also sends a text msg, often the array or
        image name.

        Arguments:
          A: numpy array or OpenCV image.
          msg: (optional) array name, image name or text message.
          flags: (optional) zmq flags.
          copy: (optional) zmq copy flag.
          track: (optional) zmq track flag.
        """

        md = dict(
            msg=msg,
            dtype=str(A.dtype),
            shape=A.shape,
        )
        self.send_json(md, flags | zmq.SNDMORE)
        return self.send(A, flags, copy=copy, track=track)

    def send_data(self,
                 msg='NoName',
                 data=b'00',
                 flags=0,
                 copy=True,
                 track=False):
        """Send a jpg buffer with a text message.

        Sends a jpg bytestring of an OpenCV image.
        Also sends text msg, often the image name.

        Arguments:
          msg: image name or text message.
          data: binary data to be sent.
          flags: (optional) zmq flags.
          copy: (optional) zmq copy flag.
          track: (optional) zmq track flag.
        """

        md = dict(msg=msg, )
        self.send_json(md, flags | zmq.SNDMORE)
        return self.send(data, flags, copy=copy, track=track)

    def recv_array(self, flags=0, copy=True, track=False):
        """Receives a numpy array with metadata and text message.

        Receives a numpy array with the metadata necessary
        for reconstructing the array (dtype,shape).
        Returns the array and a text msg, often the array or image name.

        Arguments:
          flags: (optional) zmq flags.
          copy: (optional) zmq copy flag.
          track: (optional) zmq track flag.

        Returns:
          msg: image name or text message.
          A: numpy array or OpenCV image reconstructed with dtype and shape.
        """

        md = self.recv_json(flags=flags)
        msg = self.recv(flags=flags, copy=copy, track=track)
        A = np.frombuffer(msg, dtype=md['dtype'])
        return (md['msg'], A.reshape(md['shape']))

    def recv_data(self, flags=0, copy=True, track=False):
        """Receives a jpg buffer and a text msg.

        Receives a jpg bytestring of an OpenCV image.
        Also receives a text msg, often the image name.

        Arguments:
          flags: (optional) zmq flags.
          copy: (optional) zmq copy flag.
          track: (optional) zmq track flag.

        Returns:
          msg: image name or text message.
          data: bytestring, containing data.
        """

        md = self.recv_json(flags=flags)  # metadata text
        data = self.recv(flags=flags, copy=copy, track=track)
        return (md['msg'], data)


class SerializingContext(zmq.Context):
    _socket_class = SerializingSocket



File: afy/utils.py

import sys
import time
from collections import defaultdict

import numpy as np
import cv2


def log(*args, file=sys.stderr, **kwargs):
    time_str = f'{time.time():.6f}'
    print(f'[{time_str}]', *args, file=file, **kwargs)


def info(*args, file=sys.stdout, **kwargs):
    print(*args, file=file, **kwargs)


class Tee(object):
    def __init__(self, filename, mode='w', terminal=sys.stderr):
        self.file = open(filename, mode, buffering=1)
        self.terminal = terminal

    def __del__(self):
        self.file.close()

    def write(self, *args, **kwargs):
        log(*args, file=self.file, **kwargs)
        log(*args, file=self.terminal, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.write(*args, **kwargs)

    def flush(self):
        self.file.flush()


class Logger():
    def __init__(self, filename, verbose=True):
        self.tee = Tee(filename)
        self.verbose = verbose

    def __call__(self, *args, important=False, **kwargs):
        if not self.verbose and not important:
            return

        self.tee(*args, **kwargs)


class Once():
    _id = {}

    def __init__(self, what, who=log, per=1e12):
        """ Do who(what) once per seconds.
        what: args for who
        who: callable
        per: frequency in seconds.
        """
        assert callable(who)
        now = time.time()
        if what not in Once._id or now - Once._id[what] > per:
            who(what)
            Once._id[what] = now


class TicToc:
    def __init__(self):
        self.t = None
        self.t_init = time.time()

    def tic(self):
        self.t = time.time()

    def toc(self, total=False):
        if total:
            return (time.time() - self.t_init) * 1000

        assert self.t, 'You forgot to call tic()'
        return (time.time() - self.t) * 1000

    def tocp(self, str):
        t = self.toc()
        log(f"{str} took {t:.4f}ms")
        return t


class AccumDict:
    def __init__(self, num_f=3):
        self.d = defaultdict(list)
        self.num_f = num_f
        
    def add(self, k, v):
        self.d[k] += [v]
        
    def __dict__(self):
        return self.d

    def __getitem__(self, key):
        return self.d[key]
    
    def __str__(self):
        s = ''
        for k in self.d:
            if not self.d[k]:
                continue
            cur = self.d[k][-1]
            avg = np.mean(self.d[k])
            format_str = '{:.%df}' % self.num_f
            cur_str = format_str.format(cur)
            avg_str = format_str.format(avg)
            s += f'{k} {cur_str} ({avg_str})\t\t'
        return s
    
    def __repr__(self):
        return self.__str__()


def clamp(value, min_value, max_value):
    return max(min(value, max_value), min_value)


def crop(img, p=0.7, offset_x=0, offset_y=0):
    h, w = img.shape[:2]
    x = int(min(w, h) * p)
    l = (w - x) // 2
    r = w - l
    u = (h - x) // 2
    d = h - u

    offset_x = clamp(offset_x, -l, w - r)
    offset_y = clamp(offset_y, -u, h - d)

    l += offset_x
    r += offset_x
    u += offset_y
    d += offset_y

    return img[u:d, l:r], (offset_x, offset_y)


def pad_img(img, target_size, default_pad=0):
    sh, sw = img.shape[:2]
    w, h = target_size
    pad_w, pad_h = default_pad, default_pad
    if w / h > 1:
        pad_w += int(sw * (w / h) - sw) // 2
    else:
        pad_h += int(sh * (h / w) - sh) // 2
    out = np.pad(img, [[pad_h, pad_h], [pad_w, pad_w], [0,0]], 'constant')
    return out


def resize(img, size, version='cv'):
    return cv2.resize(img, size)



File: afy/predictor_local.py

from scipy.spatial import ConvexHull
import torch
import yaml
from modules.keypoint_detector import KPDetector
from modules.generator_optim import OcclusionAwareGenerator
from sync_batchnorm import DataParallelWithCallback
import numpy as np
import face_alignment


def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,
                 use_relative_movement=False, use_relative_jacobian=False):
    if adapt_movement_scale:
        source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume
        driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume
        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
    else:
        adapt_movement_scale = 1

    kp_new = {k: v for k, v in kp_driving.items()}

    if use_relative_movement:
        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
        kp_value_diff *= adapt_movement_scale
        kp_new['value'] = kp_value_diff + kp_source['value']

        if use_relative_jacobian:
            jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))
            kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])

    return kp_new


def to_tensor(a):
    return torch.tensor(a[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2) / 255


class PredictorLocal:
    def __init__(self, config_path, checkpoint_path, relative=False, adapt_movement_scale=False, device=None, enc_downscale=1):
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        self.relative = relative
        self.adapt_movement_scale = adapt_movement_scale
        self.start_frame = None
        self.start_frame_kp = None
        self.kp_driving_initial = None
        self.config_path = config_path
        self.checkpoint_path = checkpoint_path
        self.generator, self.kp_detector = self.load_checkpoints()
        self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True, device=self.device)
        self.source = None
        self.kp_source = None
        self.enc_downscale = enc_downscale

    def load_checkpoints(self):
        with open(self.config_path) as f:
            config = yaml.load(f, Loader=yaml.FullLoader)
    
        generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
                                            **config['model_params']['common_params'])
        generator.to(self.device)
    
        kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
                                 **config['model_params']['common_params'])
        kp_detector.to(self.device)
    
        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
        generator.load_state_dict(checkpoint['generator'])
        kp_detector.load_state_dict(checkpoint['kp_detector'])
    
        generator.eval()
        kp_detector.eval()
        
        return generator, kp_detector

    def reset_frames(self):
        self.kp_driving_initial = None

    def set_source_image(self, source_image):
        self.source = to_tensor(source_image).to(self.device)
        self.kp_source = self.kp_detector(self.source)

        if self.enc_downscale > 1:
            h, w = int(self.source.shape[2] / self.enc_downscale), int(self.source.shape[3] / self.enc_downscale)
            source_enc = torch.nn.functional.interpolate(self.source, size=(h, w), mode='bilinear')
        else:
            source_enc = self.source

        self.generator.encode_source(source_enc)

    def predict(self, driving_frame):
        assert self.kp_source is not None, "call set_source_image()"
        
        with torch.no_grad():
            driving = to_tensor(driving_frame).to(self.device)

            if self.kp_driving_initial is None:
                self.kp_driving_initial = self.kp_detector(driving)
                self.start_frame = driving_frame.copy()
                self.start_frame_kp = self.get_frame_kp(driving_frame)

            kp_driving = self.kp_detector(driving)
            kp_norm = normalize_kp(kp_source=self.kp_source, kp_driving=kp_driving,
                                kp_driving_initial=self.kp_driving_initial, use_relative_movement=self.relative,
                                use_relative_jacobian=self.relative, adapt_movement_scale=self.adapt_movement_scale)

            out = self.generator(self.source, kp_source=self.kp_source, kp_driving=kp_norm)

            out = np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0]
            out = (np.clip(out, 0, 1) * 255).astype(np.uint8)

            return out

    def get_frame_kp(self, image):
        kp_landmarks = self.fa.get_landmarks(image)
        if kp_landmarks:
            kp_image = kp_landmarks[0]
            kp_image = self.normalize_alignment_kp(kp_image)
            return kp_image
        else:
            return None

    @staticmethod
    def normalize_alignment_kp(kp):
        kp = kp - kp.mean(axis=0, keepdims=True)
        area = ConvexHull(kp[:, :2]).volume
        area = np.sqrt(area)
        kp[:, :2] = kp[:, :2] / area
        return kp
    
    def get_start_frame(self):
        return self.start_frame

    def get_start_frame_kp(self):
        return self.start_frame_kp



File: afy/predictor_remote.py

from arguments import opt
from networking import SerializingContext, check_connection
from utils import Logger, TicToc, AccumDict, Once

import multiprocessing as mp
import queue

import cv2
import numpy as np
import zmq
import msgpack
import msgpack_numpy as m
m.patch()


PUT_TIMEOUT = 0.1 # s
GET_TIMEOUT = 0.1 # s
RECV_TIMEOUT = 1000 # ms
QUEUE_SIZE = 100


class PredictorRemote:
    def __init__(self, *args, in_addr=None, out_addr=None, **kwargs):
        self.in_addr = in_addr
        self.out_addr = out_addr
        self.predictor_args = (args, kwargs)
        self.timing = AccumDict()
        self.log = Logger('./var/log/predictor_remote.log', verbose=opt.verbose)

        self.send_queue = mp.Queue(QUEUE_SIZE)
        self.recv_queue = mp.Queue(QUEUE_SIZE)

        self.worker_alive = mp.Value('i', 0)

        self.send_process = mp.Process(
            target=self.send_worker, 
            args=(self.in_addr, self.send_queue, self.worker_alive),
            name="send_process"
            )
        self.recv_process = mp.Process(
            target=self.recv_worker, 
            args=(self.out_addr, self.recv_queue, self.worker_alive),
            name="recv_process"
            )

        self._i_msg = -1

    def start(self):
        self.worker_alive.value = 1
        self.send_process.start()
        self.recv_process.start()

        self.init_remote_worker()

    def stop(self):
        self.worker_alive.value = 0
        self.log("join worker processes...")
        self.send_process.join(timeout=5)
        self.recv_process.join(timeout=5)
        self.send_process.terminate()
        self.recv_process.terminate()

    def init_remote_worker(self):
        return self._send_recv_async('__init__', self.predictor_args, critical=True)

    def __getattr__(self, attr):
        is_critical = attr != 'predict'
        return lambda *args, **kwargs: self._send_recv_async(attr, (args, kwargs), critical=is_critical)

    def _send_recv_async(self, method, args, critical):
        self._i_msg += 1

        args, kwargs = args

        tt = TicToc()
        tt.tic()
        if method == 'predict':
            image = args[0]
            assert isinstance(image, np.ndarray), 'Expected image'
            ret_code, data = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), opt.jpg_quality])
        else:
            data = msgpack.packb((args, kwargs))
        self.timing.add('PACK', tt.toc())

        meta = {
            'name': method,
            'critical': critical,
            'id': self._i_msg
        }

        self.log("send", meta)

        if critical:
            self.send_queue.put((meta, data))

            while True:
                meta_recv, data_recv = self.recv_queue.get()
                if meta_recv == meta:
                    break
        else:
            try:
                # TODO: find good timeout
                self.send_queue.put((meta, data), timeout=PUT_TIMEOUT)
            except queue.Full:
                self.log('send_queue is full')

            try:
                meta_recv, data_recv = self.recv_queue.get(timeout=GET_TIMEOUT)
            except queue.Empty:
                self.log('recv_queue is empty')
                return None

        self.log("recv", meta_recv)

        tt.tic()
        if meta_recv['name'] == 'predict':
            result = cv2.imdecode(np.frombuffer(data_recv, dtype='uint8'), -1)
        else:
            result = msgpack.unpackb(data_recv)
        self.timing.add('UNPACK', tt.toc())

        if opt.verbose:
            Once(self.timing, per=1)

        return result

    @staticmethod
    def send_worker(address, send_queue, worker_alive):
        timing = AccumDict()
        log = Logger('./var/log/send_worker.log', opt.verbose)

        ctx = SerializingContext()
        sender = ctx.socket(zmq.PUSH)
        sender.connect(address)

        log(f"Sending to {address}")

        try:
            while worker_alive.value:
                tt = TicToc()

                try:
                    msg = send_queue.get(timeout=GET_TIMEOUT)
                except queue.Empty:
                    continue

                tt.tic()
                sender.send_data(*msg)
                timing.add('SEND', tt.toc())

                if opt.verbose:
                    Once(timing, log, per=1)
        except KeyboardInterrupt:
            log("send_worker: user interrupt")
        finally:
            worker_alive.value = 0

        sender.disconnect(address)
        sender.close()
        ctx.destroy()
        log("send_worker exit")

    @staticmethod
    def recv_worker(address, recv_queue, worker_alive):
        timing = AccumDict()
        log = Logger('./var/log/recv_worker.log')

        ctx = SerializingContext()
        receiver = ctx.socket(zmq.PULL)
        receiver.connect(address)
        receiver.RCVTIMEO = RECV_TIMEOUT

        log(f"Receiving from {address}")

        try:
            while worker_alive.value:
                tt = TicToc()

                try:
                    tt.tic()
                    msg = receiver.recv_data()
                    timing.add('RECV', tt.toc())
                except zmq.error.Again:
                    continue
                
                try:
                    recv_queue.put(msg, timeout=PUT_TIMEOUT)
                except queue.Full:
                    log('recv_queue full')
                    continue

                if opt.verbose:
                    Once(timing, log, per=1)
        except KeyboardInterrupt:
            log("recv_worker: user interrupt")
        finally:
            worker_alive.value = 0

        receiver.disconnect(address)
        receiver.close()
        ctx.destroy()
        log("recv_worker exit")



File: afy/camera_selector.py

import cv2
import numpy as np
import yaml

from afy.utils import log


g_selected_cam = None


def query_cameras(n_cams):
    cam_frames = {}
    cap = None
    for camid in range(n_cams):
        log(f"Trying camera with id {camid}")
        cap = cv2.VideoCapture(camid)

        if not cap.isOpened():
            log(f"Camera with id {camid} is not available")
            continue

        ret, frame = cap.read()

        if not ret or frame is None:
            log(f"Could not read from camera with id {camid}")
            cap.release()
            continue

        for i in range(10):
            ret, frame = cap.read()

        cam_frames[camid] = frame.copy()

        cap.release()

    return cam_frames


def make_grid(images, cell_size=(320, 240), cols=2):
    w0, h0 = cell_size
    _rows = len(images) // cols + int(len(images) % cols)
    _cols = min(len(images), cols)
    grid = np.zeros((h0 * _rows, w0 * _cols, 3), dtype=np.uint8)
    for i, (camid, img) in enumerate(images.items()):
        img = cv2.resize(img, (w0, h0))
        # add rect
        img = cv2.rectangle(img, (1, 1), (w0 - 1, h0 - 1), (0, 0, 255), 2)
        # add id
        img = cv2.putText(img, f'Camera {camid}', (10, 30), 0, 1, (0, 255, 0), 2)
        c = i % cols
        r = i // cols
        grid[r * h0:(r + 1) * h0, c * w0:(c + 1) * w0] = img[..., :3]
    return grid


def mouse_callback(event, x, y, flags, userdata):
    global g_selected_cam
    if event == 1:
        cell_size, grid_cols, cam_frames = userdata
        c = x // cell_size[0]
        r = y // cell_size[1]
        camid = r * grid_cols + c
        if camid < len(cam_frames):
            g_selected_cam = camid


def select_camera(cam_frames, window="Camera selector"):
    cell_size = 320, 240
    grid_cols = 2
    grid = make_grid(cam_frames, cols=grid_cols)

    # to fit the text if only one cam available
    if grid.shape[1] == 320:
        cell_size = 640, 480
        grid = cv2.resize(grid, cell_size)

    cv2.putText(grid, f'Click on the web camera to use', (10, grid.shape[0] - 30), 0, 0.7, (200, 200, 200), 2)

    cv2.namedWindow(window)
    cv2.setMouseCallback(window, mouse_callback, (cell_size, grid_cols, cam_frames))
    cv2.imshow(window, grid)

    while True:
        key = cv2.waitKey(10)

        if g_selected_cam is not None:
            break
        
        if key == 27:
            break

    cv2.destroyAllWindows()

    if g_selected_cam is not None:
        return list(cam_frames)[g_selected_cam]
    else:
        return list(cam_frames)[0]


if __name__ == '__main__':
    with open('config.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    cam_frames = query_cameras(config['query_n_cams'])

    if cam_frames:
        selected_cam = select_camera(cam_frames)
        print(f"Selected camera {selected_cam}")
    else:
        log("No cameras are available")




File: afy/predictor_worker.py

from predictor_local import PredictorLocal
from arguments import opt
from networking import SerializingContext, check_connection
from utils import Logger, TicToc, AccumDict, Once

import cv2
import numpy as np
import zmq
import msgpack
import msgpack_numpy as m
m.patch()

import queue
import multiprocessing as mp
import traceback
import time


PUT_TIMEOUT = 1 # s
GET_TIMEOUT = 1 # s
RECV_TIMEOUT = 1000 # ms
QUEUE_SIZE = 100


# class PredictorLocal():
#     def __init__(self, *args, **kwargs):
#         pass

#     def __getattr__(self, *args, **kwargs):
#         return lambda *args, **kwargs: None


class PredictorWorker():
    def __init__(self, in_port=None, out_port=None):
        self.recv_queue = mp.Queue(QUEUE_SIZE)
        self.send_queue = mp.Queue(QUEUE_SIZE)

        self.worker_alive = mp.Value('i', 0)

        self.recv_process = mp.Process(target=self.recv_worker, args=(in_port, self.recv_queue, self.worker_alive))
        self.predictor_process = mp.Process(target=self.predictor_worker, args=(self.recv_queue, self.send_queue, self.worker_alive))
        self.send_process = mp.Process(target=self.send_worker, args=(out_port, self.send_queue, self.worker_alive))
    
    def run(self):
        self.worker_alive.value = 1

        self.recv_process.start()
        self.predictor_process.start()
        self.send_process.start()

        try:
            self.recv_process.join()
            self.predictor_process.join()
            self.send_process.join()
        except KeyboardInterrupt:
            pass

    @staticmethod
    def recv_worker(port, recv_queue, worker_alive):
        timing = AccumDict()
        log = Logger('./var/log/recv_worker.log', verbose=opt.verbose)

        ctx = SerializingContext()
        socket = ctx.socket(zmq.PULL)
        socket.bind(f"tcp://*:{port}")
        socket.RCVTIMEO = RECV_TIMEOUT

        log(f'Receiving on port {port}', important=True)

        try:
            while worker_alive.value:
                tt = TicToc()

                try:
                    tt.tic()
                    msg = socket.recv_data()
                    timing.add('RECV', tt.toc())
                except zmq.error.Again:
                    log("recv timeout")
                    continue

                #log('recv', msg[0])

                method, data = msg
                if method['critical']:
                    recv_queue.put(msg)
                else:
                    try:
                        recv_queue.put(msg, block=False)
                    except queue.Full:
                        log('recv_queue full')

                Once(timing, log, per=1)
        except KeyboardInterrupt:
            log("recv_worker: user interrupt", important=True)

        worker_alive.value = 0
        log("recv_worker exit", important=True)

    @staticmethod
    def predictor_worker(recv_queue, send_queue, worker_alive):
        predictor = None
        predictor_args = ()
        timing = AccumDict()
        log = Logger('./var/log/predictor_worker.log', verbose=opt.verbose)
        
        try:
            while worker_alive.value:
                tt = TicToc()

                try:
                    method, data = recv_queue.get(timeout=GET_TIMEOUT)
                except queue.Empty:
                    continue

                # get the latest non-critical request from the queue
                # don't skip critical request
                while not recv_queue.empty() and not method['critical']:
                    log(f"skip {method}")
                    method, data = recv_queue.get()

                log("working on", method)

                try:
                    tt.tic()
                    if method['name'] == 'predict':
                        image = cv2.imdecode(np.frombuffer(data, dtype='uint8'), -1)
                    else:
                        args = msgpack.unpackb(data)
                    timing.add('UNPACK', tt.toc())
                except ValueError:
                    log("Invalid Message", important=True)
                    continue

                tt.tic()
                if method['name'] == "hello":
                    result = "OK"
                elif method['name'] == "__init__":
                    if args == predictor_args:
                        log("Same config as before... reusing previous predictor")
                    else:
                        del predictor
                        predictor_args = args
                        predictor = PredictorLocal(*predictor_args[0], **predictor_args[1])
                        log("Initialized predictor with:", predictor_args, important=True)
                    result = True
                    tt.tic() # don't account for init
                elif method['name'] == 'predict':
                    assert predictor is not None, "Predictor was not initialized"
                    result = getattr(predictor, method['name'])(image)
                else:
                    assert predictor is not None, "Predictor was not initialized"
                    result = getattr(predictor, method['name'])(*args[0], **args[1])
                timing.add('CALL', tt.toc())

                tt.tic()
                if method['name'] == 'predict':
                    assert isinstance(result, np.ndarray), f'Expected np.ndarray, got {result.__class__}'
                    ret_code, data_send = cv2.imencode(".jpg", result, [int(cv2.IMWRITE_JPEG_QUALITY), opt.jpg_quality])
                else:
                    data_send = msgpack.packb(result)
                timing.add('PACK', tt.toc())

                if method['critical']:
                    send_queue.put((method, data_send))
                else:
                    try:
                        send_queue.put((method, data_send), block=False)
                    except queue.Full:
                        log("send_queue full")
                        pass

                Once(timing, log, per=1)
        except KeyboardInterrupt:
            log("predictor_worker: user interrupt", important=True)
        except Exception as e:
            log("predictor_worker error", important=True)
            traceback.print_exc()
    
        worker_alive.value = 0
        log("predictor_worker exit", important=True)

    @staticmethod
    def send_worker(port, send_queue, worker_alive):
        timing = AccumDict()
        log = Logger('./var/log/send_worker.log', verbose=opt.verbose)

        ctx = SerializingContext()
        socket = ctx.socket(zmq.PUSH)
        socket.bind(f"tcp://*:{port}")

        log(f'Sending on port {port}', important=True)

        try:
            while worker_alive.value:
                tt = TicToc()

                try:
                    method, data = send_queue.get(timeout=GET_TIMEOUT)
                except queue.Empty:
                    log("send queue empty")
                    continue

                # get the latest non-critical request from the queue
                # don't skip critical request
                while not send_queue.empty() and not method['critical']:
                    log(f"skip {method}")
                    method, data = send_queue.get()

                log("sending", method)

                tt.tic()
                socket.send_data(method, data)
                timing.add('SEND', tt.toc())

                Once(timing, log, per=1)
        except KeyboardInterrupt:
            log("predictor_worker: user interrupt", important=True)

        worker_alive.value = 0
        log("send_worker exit", important=True)


def run_worker(in_port=None, out_port=None):
    worker = PredictorWorker(in_port=in_port, out_port=out_port)
    worker.run()



![](docs/mona.gif)

# Avatarify Python

Photorealistic avatars for video-conferencing.

Avatarify Python requires manually downloading and installing some dependencies, and is therefore best suited for users who have some experience with command line applications. [Avatarify Desktop](https://github.com/alievk/avatarify-desktop), which aims to be easier to install and use, is recommended for most users. If you still want to use Avatarify Python, proceed to the [install instructions](docs/).

Based on [First Order Motion Model](https://github.com/AliaksandrSiarohin/first-order-model).

## News
- **7 Mar 2021.** Renamed project to Avatarify Python to distinguish it from other versions of Avatarify
- **14 December 2020.** Released Avatarify Desktop. Check it out [here](https://github.com/alievk/avatarify-desktop).
- **11 July 2020.** Added Docker support. Now you can run Avatarify from Docker on [Linux](https://github.com/alievk/avatarify-python/blob/master/docs/README.md#docker). Thanks to [mikaelhg](https://github.com/mikaelhg) and [mintmaker](https://github.com/mintmaker) for contribution!
- **22 May 2020.** Added [Google Colab](https://colab.research.google.com/github/alievk/avatarify/blob/master/avatarify.ipynb) mode. Now you can run Avatarify on any computer without GPU!
- **7 May 2020.** Added remote GPU support for all platforms (based on [mynameisfiber's](https://github.com/mynameisfiber) solution). [Demo](https://youtu.be/3Dz_bUIPYFM). Deployment [instructions](https://github.com/alievk/avatarify-python/wiki/Remote-GPU). 
- **24 April 2020.** Added Windows installation [tutorial](https://www.youtube.com/watch?v=lym9ANVb120).
- **17 April 2020.** Created Slack community. Please join via [invitation link](https://join.slack.com/t/avatarify/shared_invite/zt-dyoqy8tc-~4U2ObQ6WoxuwSaWKKVOgg).
- **15 April 2020.** Added [StyleGAN-generated](https://www.thispersondoesnotexist.com) avatars. Just press `Q` and now you drive a person that never existed. Every time you push the button – new avatar is sampled.
- **13 April 2020.** Added Windows support (kudos to [9of9](https://github.com/9of9)).

## Avatarify apps

We have deployed Avatarify on iOS and Android devices using our proprietary inference engine. The iOS version features the Life mode for recording animations in real time. However, the Life mode is not available on Android devices due to the diversity of the devices we have to support.

[<img src=docs/appstore-badge.png alt="drawing" height="40"/>](https://apps.apple.com/app/apple-store/id1512669147?pt=121960189&ct=GitHub&mt=8)
[<img src=docs/google-play-badge.png alt="drawing" height="40"/>](https://play.google.com/store/apps/details?id=com.avatarify.android)
