tst_scene_render.py 3.86 KB
Newer Older
wester committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
#!/usr/bin/env python


# Python 2/3 compatibility
from __future__ import print_function

import numpy as np
from numpy import pi, sin, cos

import cv2

defaultSize = 512

class TestSceneRender():

wester committed
16
    def __init__(self, bgImg = None, fgImg = None, deformation = False, noise = 0.0, speed = 0.25, **params):
wester committed
17 18 19 20
        self.time = 0.0
        self.timeStep = 1.0 / 30.0
        self.foreground = fgImg
        self.deformation = deformation
wester committed
21
        self.noise = noise
wester committed
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
        self.speed = speed

        if bgImg is not None:
            self.sceneBg = bgImg.copy()
        else:
            self.sceneBg = np.zeros(defaultSize, defaultSize, np.uint8)

        self.w = self.sceneBg.shape[0]
        self.h = self.sceneBg.shape[1]

        if fgImg is not None:
            self.foreground = fgImg.copy()
            self.center = self.currentCenter = (int(self.w/2 - fgImg.shape[0]/2), int(self.h/2 - fgImg.shape[1]/2))

            self.xAmpl = self.sceneBg.shape[0] - (self.center[0] + fgImg.shape[0])
            self.yAmpl = self.sceneBg.shape[1] - (self.center[1] + fgImg.shape[1])

        self.initialRect = np.array([ (self.h/2, self.w/2), (self.h/2, self.w/2 + self.w/10),
         (self.h/2 + self.h/10, self.w/2 + self.w/10), (self.h/2 + self.h/10, self.w/2)]).astype(int)
        self.currentRect = self.initialRect
wester committed
42
        np.random.seed(10)
wester committed
43 44

    def getXOffset(self, time):
wester committed
45
        return int(self.xAmpl*cos(time*self.speed))
wester committed
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87


    def getYOffset(self, time):
        return int(self.yAmpl*sin(time*self.speed))

    def setInitialRect(self, rect):
        self.initialRect = rect

    def getRectInTime(self, time):

        if self.foreground is not None:
            tmp = np.array(self.center) + np.array((self.getXOffset(time), self.getYOffset(time)))
            x0, y0 = tmp
            x1, y1 = tmp + self.foreground.shape[0:2]
            return np.array([y0, x0, y1, x1])
        else:
            x0, y0 = self.initialRect[0] + np.array((self.getXOffset(time), self.getYOffset(time)))
            x1, y1 = self.initialRect[2] + np.array((self.getXOffset(time), self.getYOffset(time)))
            return np.array([y0, x0, y1, x1])

    def getCurrentRect(self):

        if self.foreground is not None:

            x0 = self.currentCenter[0]
            y0 = self.currentCenter[1]
            x1 = self.currentCenter[0] + self.foreground.shape[0]
            y1 = self.currentCenter[1] + self.foreground.shape[1]
            return np.array([y0, x0, y1, x1])
        else:
            x0, y0 = self.currentRect[0]
            x1, y1 = self.currentRect[2]
            return np.array([x0, y0, x1, y1])

    def getNextFrame(self):
        img = self.sceneBg.copy()

        if self.foreground is not None:
            self.currentCenter = (self.center[0] + self.getXOffset(self.time), self.center[1] + self.getYOffset(self.time))
            img[self.currentCenter[0]:self.currentCenter[0]+self.foreground.shape[0],
             self.currentCenter[1]:self.currentCenter[1]+self.foreground.shape[1]] = self.foreground
        else:
wester committed
88
            self.currentRect = self.initialRect + np.int( 30*cos(self.time) + 50*sin(self.time/3))
wester committed
89
            if self.deformation:
wester committed
90
                self.currentRect[1:3] += int(self.h/20*cos(self.time))
wester committed
91 92 93
            cv2.fillConvexPoly(img, self.currentRect, (0, 0, 255))

        self.time += self.timeStep
wester committed
94 95 96 97 98

        if self.noise:
            noise = np.zeros(self.sceneBg.shape, np.int8)
            cv2.randn(noise, np.zeros(3), np.ones(3)*255*self.noise)
            img = cv2.add(img, noise, dtype=cv2.CV_8UC3)
wester committed
99 100 101 102 103 104 105 106
        return img

    def resetTime(self):
        self.time = 0.0


if __name__ == '__main__':

wester committed
107
    backGr = cv2.imread('../../../samples/data/lena.jpg')
wester committed
108

wester committed
109
    render = TestSceneRender(backGr, noise = 0.5)
wester committed
110 111 112 113 114 115 116

    while True:

        img = render.getNextFrame()
        cv2.imshow('img', img)

        ch = cv2.waitKey(3)
wester committed
117
        if ch == 27:
wester committed
118 119
            break
    cv2.destroyAllWindows()