lib/pymaf/utils/streamer.py

import cv2
import torch
import numpy as np
import imageio


def aug_matrix(w1, h1, w2, h2):
    dx = (w2 - w1) / 2.0
    dy = (h2 - h1) / 2.0

    matrix_trans = np.array([[1.0, 0, dx], [0, 1.0, dy], [0, 0, 1.0]])

    scale = np.min([float(w2) / w1, float(h2) / h1])

    M = get_affine_matrix(center=(w2 / 2.0, h2 / 2.0),
                          translate=(0, 0),
                          scale=scale)

    M = np.array(M + [0., 0., 1.]).reshape(3, 3)
    M = M.dot(matrix_trans)

    return M


def get_affine_matrix(center, translate, scale):
    cx, cy = center
    tx, ty = translate

    M = [1, 0, 0, 0, 1, 0]
    M = [x * scale for x in M]

    # Apply translation and of center translation: RSS * C^-1
    M[2] += M[0] * (-cx) + M[1] * (-cy)
    M[5] += M[3] * (-cx) + M[4] * (-cy)

    # Apply center translation: T * C * RSS * C^-1
    M[2] += cx + tx
    M[5] += cy + ty
    return M


class BaseStreamer():
    """This streamer will return images at 512x512 size.
    """

    def __init__(self,
                 width=512,
                 height=512,
                 pad=True,
                 mean=(0.5, 0.5, 0.5),
                 std=(0.5, 0.5, 0.5),
                 **kwargs):
        self.width = width
        self.height = height
        self.pad = pad
        self.mean = np.array(mean)
        self.std = np.array(std)

        self.loader = self.create_loader()

    def create_loader(self):
        raise NotImplementedError
        yield np.zeros((600, 400, 3))  # in RGB (0, 255)

    def __getitem__(self, index):
        image = next(self.loader)
        in_height, in_width, _ = image.shape
        M = aug_matrix(in_width, in_height, self.width, self.height, self.pad)
        image = cv2.warpAffine(image,
                               M[0:2, :], (self.width, self.height),
                               flags=cv2.INTER_CUBIC)

        input = np.float32(image)
        input = (input / 255.0 - self.mean) / self.std  # TO [-1.0, 1.0]
        input = input.transpose(2, 0, 1)  # TO [3 x H x W]
        return torch.from_numpy(input).float()

    def __len__(self):
        raise NotImplementedError


class CaptureStreamer(BaseStreamer):
    """This streamer takes webcam as input.
    """

    def __init__(self, id=0, width=512, height=512, pad=True, **kwargs):
        super().__init__(width, height, pad, **kwargs)
        self.capture = cv2.VideoCapture(id)

    def create_loader(self):
        while True:
            _, image = self.capture.read()
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # RGB
            yield image

    def __len__(self):
        return 100_000_000

    def __del__(self):
        self.capture.release()


class VideoListStreamer(BaseStreamer):
    """This streamer takes a list of video files as input.
    """

    def __init__(self, files, width=512, height=512, pad=True, **kwargs):
        super().__init__(width, height, pad, **kwargs)
        self.files = files
        self.captures = [imageio.get_reader(f) for f in files]
        self.nframes = sum([
            int(cap._meta["fps"] * cap._meta["duration"])
            for cap in self.captures
        ])

    def create_loader(self):
        for capture in self.captures:
            for image in capture:  # RGB
                yield image

    def __len__(self):
        return self.nframes

    def __del__(self):
        for capture in self.captures:
            capture.close()


class ImageListStreamer(BaseStreamer):
    """This streamer takes a list of image files as input.
    """

    def __init__(self, files, width=512, height=512, pad=True, **kwargs):
        super().__init__(width, height, pad, **kwargs)
        self.files = files

    def create_loader(self):
        for f in self.files:
            image = cv2.imread(f, cv2.IMREAD_UNCHANGED)[:, :, 0:3]
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # RGB
            yield image

    def __len__(self):
        return len(self.files)