In [1]:

    

import numpy as np
import cv2
from matplotlib import pyplot as plt
import os
from scipy.misc import imsave


    

In [2]:

    

!mkdir videos
!mkdir annotations
!mkdir tmp
!mkdir tmp/rpn_tmp


    

    




mkdir: cannot create directory ‘tmp’: File exists
mkdir: cannot create directory ‘tmp/rpn_tmp’: File exists

In [3]:

    

imsize = 32


    

In [4]:

    

class Agent:
    def __init__(self, positive):
        
        self.positive = positive
        
        if positive:
            self.color = np.array([0,1,0], dtype='int64')
            self.size = np.random.randint(low=1, high=3, size=2)
        else:
            self.color = np.array([0,0,1], dtype='int64')#np.random.normal(size=3)
            self.size = np.random.randint(low=4, high=8, size=2)
        
        self.pos = np.random.randint(low=0, high=imsize-1, size=2)
        self.v = np.random.normal(size=2)*2.5
        
    def correct_pv(self):
        #correct position and velocity
        pos = self.pos
        v = self.v
        s = self.size
        
        if pos[0] + s[0] >= imsize:
            v[0] *= -1
            pos[0] = imsize - s[0] - 1
            
        if pos[1] + s[1] >= imsize:
            v[1] *= -1
            pos[1] = imsize - s[1] - 1
        
        if pos[0] - s[0] < 0:
            v[0] *= -1
            pos[0] = s[0]
            
        if pos[1] - s[1] < 0:
            v[1] *= -1
            pos[1] = s[1]
        
        self.pos = pos
        self.v = v
    def plot_on_image(self, im):
        pos = self.pos
        s = self.size
        
        im[pos[0] - s[0]:pos[0] + s[0], pos[1] - s[1]:pos[1] + s[1]] = self.color
        
        bcol = np.array([1,1,1], dtype='float')
        im[0,:] = bcol
        im[-1,:] = bcol
        im[:,0] = bcol
        im[:,-1] = bcol
        
    def step(self):
        self.v += np.random.normal(size=2)/2
        self.v = np.clip(self.v, -3, 3)
        
        self.pos = (self.v + self.pos).astype('int64')
        self.correct_pv()
    
    def bbox(self):
        pos = self.pos
        s = self.size
        return [pos[1] - s[1], pos[0] - s[0], pos[1] + s[1], pos[0] + s[0]]


    

In [5]:

    

def create_video(num_positive, num_negative, length):
    im = np.zeros([imsize, imsize, 3])
    apos = [Agent(True) for _ in range(num_positive)]
    afal = [Agent(False) for _ in range(num_negative)]
    
    imgs = []
    bboxes = []
    
    for _ in range(length):
        im = np.zeros([imsize, imsize, 3])

        for agents in [apos, afal]:
            [ag.step() for ag in agents]
            [ag.plot_on_image(im) for ag in agents]

        imgs.append(im)
        
        bboxes.append( [ag.bbox() for ag in apos] )
        
    return np.array(imgs), bboxes


    

In [6]:

    

video, bboxes = create_video(6,3,4)


    

In [7]:

    

video.min()


    

    
Out[7]:





0.0

In [8]:

    

for im, bb in zip(video, bboxes):
    for b in bb:
        cv2.rectangle(im,(b[0], b[1]), (b[2], b[3]), 70)
        pass
    plt.imshow(im)
    plt.show()


    

In [9]:

    

bboxes[0]


    

    
Out[9]:





[[2, 8, 4, 10],
 [19, 14, 23, 16],
 [14, 17, 16, 21],
 [14, 10, 16, 12],
 [6, 1, 8, 5],
 [21, 20, 25, 24]]

In [10]:

    

video_dir = './videos/'
annotations_dir = './annotations/'


    

In [11]:

    

!rm -R videos/*
!rm annotations/*


    

    




rm: cannot remove 'videos/*': No such file or directory
rm: cannot remove 'annotations/*': No such file or directory

In [12]:

    

for vnum in range(3):
    cdir = os.path.join(video_dir,str(vnum))
    os.mkdir(cdir)
    
    num_pos = np.random.randint(1, 6)
    
    video, frame_bbs = create_video(num_pos,2,np.random.randint(low=50, high=100))
    for i, frame in enumerate(video):
        imsave(os.path.join(cdir, str(i) + '.jpg'), frame)
    
    annot_file = os.path.join(annotations_dir,str(vnum)) + '.txt'
    with open(annot_file, 'w') as f:
        for i, frame in enumerate(frame_bbs):
            for bb in frame:
                f.write(','.join([str(i)] + list(map(str, bb))) + '\n')


    

In [ ]: