import os, sys, glob
import csv
import cv2
import yaml
import numpy as np

import time


def main(argv):
    if(len(sys.argv) != 3):
        print('Usage ./transforms.py [root_dir] [config.yml]')
        sys.exit(1)
    #Load data
    root_dir = argv[1]  + '/' #in case we forgot
    print('Opening config file' + root_dir + argv[2])
    with open(root_dir + argv[2], 'r') as configfile:
        config = yaml.safe_load(configfile)

    np.set_printoptions(suppress=True)
    data_dir = config['movie_dir']
    video_name_regex = data_dir + "/*.avi"
    weights_dir = root_dir + config['weights_dir']
    your_weights = weights_dir + config['specific_weights']
    generic_weights = weights_dir + config['generic_weights']
    trained_weights = weights_dir + config['trained_weights']
    #train_images =  glob.glob( image_dir + "*.png" )
    #video_file1 = 'out.avi'

    #TODO allow different resolution
    width = 3840#1920
    height = 2176#1080

    display = config['display']
    showDetections = config['showDetections']

    filelist = glob.glob(video_name_regex)

    for input_file in filelist:

        direct, ext = os.path.split(input_file)
        noext, _ = os.path.splitext(ext)
        data_file = data_dir + '/tracks/' +  noext + '_POS.txt'
        video_file = data_dir + '/tracks/' +  noext + '_TR.avi'
        tr_file = data_dir + '/tracks/' +  noext + '_MAT.npy'
        if os.path.isfile(data_file):
            continue
        if os.path.isfile(tr_file):
            continue
        print(input_file, video_file)

        ##########################################################################
        ##          open the video file for inputs 
        ##########################################################################
        cap = cv2.VideoCapture(input_file)
        fps = round(cap.get(cv2.CAP_PROP_FPS))
        nframes = round(cap.get(cv2.CAP_PROP_FRAME_COUNT))


        ##########################################################################
        ##          corrections for camera motion
        ##########################################################################
        warp_mode = cv2.MOTION_AFFINE
        number_of_iterations = 20
        # Specify the threshold of the increment in the correlation coefficient between two iterations
        termination_eps = 1e-6;
        # Define termination criteria
        criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations,  termination_eps)

        im1_gray = np.array([])
        warp_matrix = np.eye(3, 3, dtype=np.float32) 
        full_warp = np.eye(3, 3, dtype=np.float32)

        save_warp = np.zeros((nframes,3,3,))
        np.save(tr_file,save_warp)



        frame_idx=0
        for i in range(nframes):

            ret, frame = cap.read() 
            sys.stdout.write('\r')
            sys.stdout.write("[%-20s] %d%% %d/%d" % ('='*int(20*i/float(nframes)), int(100.0*i/float(nframes)), i,nframes)) 
            sys.stdout.flush()

            if not(im1_gray.size):
                # enhance contrast in the image
                im1_gray = cv2.equalizeHist(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY))#[200:-200,200:-200])

            im2_gray =  cv2.equalizeHist(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY))#[200:-200,200:-200])


 #           start=time.time()
            try:
                # find difference in movement between this frame and the last frame
                (cc, warp_matrix) = cv2.findTransformECC(im1_gray,im2_gray,warp_matrix, warp_mode, criteria)
                # this frame becames the last frame for the next iteration
                im1_gray = im2_gray.copy()
            except cv2.error as e:
                warp_matrix = np.eye(3, 3, dtype=np.float32)

            # all moves are accumulated into a matrix
            full_warp = np.dot(warp_matrix,full_warp)
            save_warp[i,:,:] = full_warp
  #          print('ecc ', time.time()-start)

        np.save(tr_file,save_warp)



if __name__ == '__main__':
    main(sys.argv)

@@ -1068,16 +1068,35 @@
 regex)%0A%0A
+    scalefact=4.0%0A%0A
     for 
@@ -1403,67 +1403,8 @@
 py'%0A
-        if os.path.isfile(data_file):%0A            continue%0A
     
@@ -2115,14 +2115,17 @@
 ION_
-AFFINE
+EUCLIDEAN
 %0A   
@@ -2491,33 +2491,33 @@
 matrix = np.eye(
-3
+2
 , 3, dtype=np.fl
@@ -2942,16 +2942,95 @@
 lush()%0A%0A
+            frame2 =cv2.resize(frame,None, fx=1.0/scalefact, fy=1.0/scalefact)%0A
         
@@ -3158,32 +3158,33 @@
 2.cvtColor(frame
+2
 ,cv2.COLOR_BGR2G
@@ -3270,16 +3270,17 @@
 or(frame
+2
 ,cv2.COL
@@ -3739,33 +3739,33 @@
 matrix = np.eye(
-3
+2
 , 3, dtype=np.fl
@@ -3764,32 +3764,88 @@
 pe=np.float32)%0A%0A
+            warp_matrix%5B:,2%5D=warp_matrix%5B:,2%5D*scalefact%0A
             # al
@@ -3913,16 +3913,37 @@
  np.dot(
+full_warp,np.vstack((
 warp_mat
@@ -3946,25 +3946,25 @@
 _matrix,
-full_warp
+%5B0,0,1%5D))
 )%0A      

fix ECC