gen-rand-train-data.py 2.92 KB
Newer Older
Gihan Jayatilaka's avatar
a  
Gihan Jayatilaka committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
import sys
import numpy as np
import cv2


def makeVideo(READ_VIDEO_FILE,msg,WRITE_VIDEO_FILE,DELTA):

    GRID_Y=msg.shape[1]
    GRID_X=msg.shape[2]

    inputVideo = cv2.VideoCapture(READ_VIDEO_FILE)
    NO_FRAMES_MAX=min(msg.shape[0],inputVideo.get(cv2.CAP_PROP_FRAME_COUNT))-1
    ret,frame=inputVideo.read()
    WIDTH=frame.shape[1]
    HEIGHT=frame.shape[0]
    inputVideo.release()
    inputVideo = cv2.VideoCapture(READ_VIDEO_FILE)

    video = cv2.VideoWriter(WRITE_VIDEO_FILE,cv2.VideoWriter_fourcc('M','J','P','G'),20,(WIDTH,HEIGHT))

    one = np.ones((HEIGHT, WIDTH, 1), np.uint8)*255
    zero = np.zeros((HEIGHT, WIDTH, 1), np.uint8)
    white = cv2.merge([one*0.6, one*0.6, one*0.6]).astype(np.uint8)
    red = cv2.merge([zero, zero, one])
    green = cv2.merge([zero, one, zero])
    blue = cv2.merge([one, zero, zero])


    for col in [red,green,blue,red,green,blue,white,white,white,white]:
        video.write(col)

    frameIndex=0
    while inputVideo.isOpened():
        if frameIndex==NO_FRAMES_MAX:
            break
        ret, frame = inputVideo.read()

        if ret:
            mask = msg[frameIndex].astype(np.float32)

            mask = np.repeat(mask, HEIGHT/GRID_Y, axis=0)
            mask = np.repeat(mask, WIDTH/GRID_X, axis=1)


            raw = cv2.cvtColor(cv2.resize(frame, (WIDTH, HEIGHT)), cv2.COLOR_BGR2LAB)
            l,a,b = cv2.split(raw)
            lnorm = np.square(l*100.0/255-50)
            upper = lnorm*0.015
            lower = np.sqrt(lnorm + 20)
            delt = np.divide((upper + lower), lower)*DELTA[frameIndex]*2.55

            lmod = l + np.multiply(mask, delt)
            calib = np.max(lmod)

            if calib > 255:
                lmod = np.uint8((lmod/calib)*255)
            else:
                lmod = np.uint8(lmod)
            img = cv2.merge([lmod, a, b])
            img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
            video.write(img)

        else:
            break
        frameIndex+=1

        if frameIndex%100==0:
            print("{} frames of video encoded".format(frameIndex))


    for col in [red,green,blue,red,green,blue,white]:
        video.write(col)

    inputVideo.release()
    video.release()
    cv2.destroyAllWindows()






if __name__ == "__main__":
    np.random.seed(0)
    print("python gen-rand-train-data.py video.mp4 10000 10  10  1.0  10.0 vidOut.mp4")
    print("python gen-rand-train-data.py [1]       [2]   [3] [4] [5]  [6]   [7]")

    VIDEO_IN_FILE=sys.argv[1]
    NO_FRAMES=int(sys.argv[2])
    GRID_HEIGHT=int(sys.argv[3])
    GRID_WIDTH=int(sys.argv[4])
    DELTA_MIN=float(sys.argv[5])
    DELTA_MAX=float(sys.argv[6])
    VIDEO_OUT_FILE=sys.argv[7]

    randMsg=np.random.randint(low=0,high=2,size=(NO_FRAMES,GRID_HEIGHT,GRID_WIDTH),dtype=np.int32)
    deltas=np.random.uniform(low=DELTA_MIN,high=DELTA_MAX,size=(NO_FRAMES))


    print(deltas)
    makeVideo(VIDEO_IN_FILE,randMsg,VIDEO_OUT_FILE,deltas)






    print("FINISHED SUCCESSFULLY")