learn_prior.py 4.61 KB
Newer Older
Vladislav Samsonov's avatar
Vladislav Samsonov committed
1 2 3 4
#!/usr/bin/env python

import os
import sys
5
import numpy as np
6
import cv2 as cv
7
import struct
Vladislav Samsonov's avatar
Vladislav Samsonov committed
8
import argparse
9 10
from math import sqrt

Vladislav Samsonov's avatar
Vladislav Samsonov committed
11 12 13 14 15 16 17 18
argparser = argparse.ArgumentParser(
    description='''Use this script to generate prior for using with PCAFlow.
Basis size here must match corresponding parameter in the PCAFlow.
Gamma should be selected experimentally.''')

argparser.add_argument('-f',
                       '--files',
                       nargs='+',
Vladislav Samsonov's avatar
Vladislav Samsonov committed
19
                       help='List of optical flow .flo files for learning. You can pass a directory here and it will be scanned recursively for .flo files.',
Vladislav Samsonov's avatar
Vladislav Samsonov committed
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
                       required=True)
argparser.add_argument('-o',
                       '--output',
                       help='Output file for prior',
                       required=True)
argparser.add_argument('--width',
                       type=int,
                       help='Size of the basis first dimension',
                       required=True,
                       default=18)
argparser.add_argument('--height',
                       type=int,
                       help='Size of the basis second dimension',
                       required=True,
                       default=14)
argparser.add_argument(
    '-g',
    '--gamma',
    type=float,
    help='Amount of regularization. The greater this parameter, the bigger will be an impact of the regularization.',
    required=True)
args = argparser.parse_args()

basis_size = (args.height, args.width)
gamma = args.gamma

46 47

def find_flo(pp):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
48 49 50 51 52 53 54 55 56 57
    f = []
    for p in pp:
        if os.path.isfile(p):
            f.append(p)
        else:
            for root, subdirs, files in os.walk(p):
                f += map(lambda x: os.path.join(root, x),
                         filter(lambda x: x.split('.')[-1] == 'flo', files))
    return list(set(f))

58 59

def load_flo(flo):
60
    with open(flo, 'rb') as f:
Vladislav Samsonov's avatar
Vladislav Samsonov committed
61 62 63 64 65 66 67 68 69 70 71 72
        magic = np.fromfile(f, np.float32, count=1)[0]
        if 202021.25 != magic:
            print('Magic number incorrect. Invalid .flo file')
        else:
            w = np.fromfile(f, np.int32, count=1)[0]
            h = np.fromfile(f, np.int32, count=1)[0]
            print('Reading %dx%d flo file %s' % (w, h, flo))
            data = np.fromfile(f, np.float32, count=2 * w * h)
            # Reshape data into 3D array (columns, rows, bands)
            flow = np.reshape(data, (h, w, 2))
            return flow[:, :, 0], flow[:, :, 1]

73 74

def get_w(m):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
75
    s = m.shape
76
    w = cv.dct(m)
Vladislav Samsonov's avatar
Vladislav Samsonov committed
77 78 79 80 81 82 83
    w *= 2.0 / sqrt(s[0] * s[1])
    #w[0,0] *= 0.5
    w[:, 0] *= sqrt(0.5)
    w[0, :] *= sqrt(0.5)
    w = w[0:basis_size[0], 0:basis_size[1]].transpose().flatten()
    return w

84 85 86 87

w1 = []
w2 = []

Vladislav Samsonov's avatar
Vladislav Samsonov committed
88 89 90 91
for flo in find_flo(args.files):
    x, y = load_flo(flo)
    w1.append(get_w(x))
    w2.append(get_w(y))
92 93 94 95 96

w1mean = sum(w1) / len(w1)
w2mean = sum(w2) / len(w2)

for i in xrange(len(w1)):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
97
    w1[i] -= w1mean
98
for i in xrange(len(w2)):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
99
    w2[i] -= w2mean
100

Vladislav Samsonov's avatar
Vladislav Samsonov committed
101 102 103 104
Q1 = sum([w1[i].reshape(-1, 1).dot(w1[i].reshape(1, -1))
          for i in xrange(len(w1))]) / len(w1)
Q2 = sum([w2[i].reshape(-1, 1).dot(w2[i].reshape(1, -1))
          for i in xrange(len(w2))]) / len(w2)
105 106 107 108
Q1 = np.matrix(Q1)
Q2 = np.matrix(Q2)

if len(w1) > 1:
Vladislav Samsonov's avatar
Vladislav Samsonov committed
109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131
    while True:
        try:
            L1 = np.linalg.cholesky(Q1)
            break
        except np.linalg.linalg.LinAlgError:
            mev = min(np.linalg.eig(Q1)[0]).real
            assert (mev < 0)
            print('Q1', mev)
            if -mev < 1e-6:
                mev = -1e-6
            Q1 += (-mev * 1.000001) * np.identity(Q1.shape[0])

    while True:
        try:
            L2 = np.linalg.cholesky(Q2)
            break
        except np.linalg.linalg.LinAlgError:
            mev = min(np.linalg.eig(Q2)[0]).real
            assert (mev < 0)
            print('Q2', mev)
            if -mev < 1e-6:
                mev = -1e-6
            Q2 += (-mev * 1.000001) * np.identity(Q2.shape[0])
132
else:
Vladislav Samsonov's avatar
Vladislav Samsonov committed
133 134
    L1 = np.identity(Q1.shape[0])
    L2 = np.identity(Q2.shape[0])
135

Vladislav Samsonov's avatar
Vladislav Samsonov committed
136 137
L1 = np.linalg.inv(L1) * gamma
L2 = np.linalg.inv(L2) * gamma
138

Vladislav Samsonov's avatar
Vladislav Samsonov committed
139 140
assert (L1.shape == L2.shape)
assert (L1.shape[0] == L1.shape[1])
141

Vladislav Samsonov's avatar
Vladislav Samsonov committed
142
f = open(args.output, 'wb')
143 144 145 146 147

f.write(struct.pack('I', L1.shape[0]))
f.write(struct.pack('I', L1.shape[1]))

for i in xrange(L1.shape[0]):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
148 149
    for j in xrange(L1.shape[1]):
        f.write(struct.pack('f', L1[i, j]))
150 151

for i in xrange(L2.shape[0]):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
152 153
    for j in xrange(L2.shape[1]):
        f.write(struct.pack('f', L2[i, j]))
154

Vladislav Samsonov's avatar
Vladislav Samsonov committed
155 156
b1 = L1.dot(w1mean.reshape(-1, 1))
b2 = L2.dot(w2mean.reshape(-1, 1))
157

Vladislav Samsonov's avatar
Vladislav Samsonov committed
158
assert (L1.shape[0] == b1.shape[0])
159 160

for i in xrange(b1.shape[0]):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
161
    f.write(struct.pack('f', b1[i, 0]))
162 163

for i in xrange(b2.shape[0]):
Vladislav Samsonov's avatar
Vladislav Samsonov committed
164
    f.write(struct.pack('f', b2[i, 0]))
165

Vladislav Samsonov's avatar
Vladislav Samsonov committed
166
f.close()