0
|
1 cimport c_numpy
|
|
2 cimport c_python
|
|
3 import numpy, sys
|
|
4
|
|
5 c_numpy.import_array()
|
|
6
|
|
7 cdef extern from 'math.h':
|
|
8 cdef double fabs(double)
|
|
9
|
|
10 cdef extern from 'sys/param.h':
|
|
11 cdef int MIN(int, int)
|
|
12 cdef int MAX(int, int)
|
|
13
|
|
14 cdef double calc_sad(c_numpy.ndarray ablock, int rs, int cs, c_numpy.ndarray refblock):
|
|
15 cdef double sad
|
|
16 cdef int i, j
|
|
17 cdef char *p_a, *p_ref
|
|
18
|
|
19 sad = 0.0
|
|
20 for i from 0 <= i < refblock.dimensions[0]:
|
|
21 p_a = ablock.data + (rs+i) * ablock.strides[0] + cs * ablock.strides[1]
|
|
22 p_ref = refblock.data + i * refblock.strides[0]
|
|
23
|
|
24 for j from 0 <= j < refblock.dimensions[1]:
|
|
25 sad = sad + fabs((<double*>p_a)[0] - (<double*>p_ref)[0])
|
|
26 p_a = p_a + ablock.strides[1]
|
|
27 p_ref = p_ref + refblock.strides[1]
|
|
28
|
|
29 return sad
|
|
30
|
|
31 def me(c_numpy.ndarray a, c_numpy.ndarray refblock, int rc, int cc, int sr):
|
|
32 cdef int rs, rs1, rs2, cs, cs1, cs2
|
|
33 cdef int min_r, min_c
|
|
34 cdef double sad, min_sad
|
|
35 cdef int bs
|
|
36
|
|
37 bs = refblock.dimensions[0]
|
|
38 min_sad = sys.maxint
|
|
39 min_r, min_c = 0, 0
|
|
40
|
|
41 rs1 = MAX(0,rc-sr)
|
|
42 rs2 = MIN(a.dimensions[0]-bs,rc+sr)+1
|
|
43
|
|
44 cs1 = MAX(0,cc-sr)
|
|
45 cs2 = MIN(cc+sr,a.dimensions[1]-bs)+1
|
|
46
|
|
47 for rs from rs1 <= rs < rs2:
|
|
48 for cs from cs1 <= cs < cs2:
|
|
49 sad = calc_sad(a, rs, cs, refblock)
|
|
50 if sad < min_sad:
|
|
51 # found new local block SAD minimum, store motion vector
|
|
52 min_r, min_c, min_sad = rs, cs, sad
|
|
53
|
|
54 return min_r, min_c, min_sad
|
|
55
|