import argparse
import sys
import os
import glob
import time
import cv2
import numpy as np
import shutil


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(
                        description='Extract homography from a list of videos')
    parser.add_argument('--video_list', dest='video_list_file',
                        help='a file that contains all videos (one video per line)',
                        default=None, type=str)
    parser.add_argument('--output_folder', dest='output_folder',
                        help='output folder to save selected video frames (root of the dataset)',
                        default=None, type=str)
 
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    return args

def image_quality(img):
    """
    Implementation of the paper
    "No-Reference Perceptual Quality Assessment of JPEG Compressed Images"
    Used as a quick filter of frames
    """
    # convert bgr image to gray -> float32
    score = 0.0
    if img is None:
        return score

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    x = gray.astype(np.float32)
    h, w = x.shape[0], x.shape[1]

    # horizontal
    d_h = x[:,1:] - x[:,:-1]
    w_bound = int(8*(np.floor(w/8.0)-1)) + 1
    B_h = np.mean(np.abs(d_h[:,7:w_bound:8]))
    A_h = (8.0 * np.mean(np.abs(d_h)) - B_h) / 7.0
    sig_h = np.sign(d_h)
    left_sig, right_sig = sig_h[:,:-2], sig_h[:,1:-1]
    Z_h = np.mean((left_sig * right_sig)<0)

    # vertical
    d_v = x[1:, :] - x[:-1, :]
    h_bound = int(8*(np.floor(h/8.0)-1)) + 1
    B_v = np.mean(np.abs(d_v[7:h_bound:8, :]))
    A_v = (8.0 * np.mean(np.abs(d_v)) - B_v) / 7.0
    sig_v = np.sign(d_v)
    up_sig, down_sig = sig_v[:-2, :], sig_v[1:-1, :]
    Z_v = np.mean((up_sig * down_sig)<0)

    # combine the weights
    B = (B_h + B_v)/2.0
    A = (A_h + A_v)/2.0
    Z = (Z_h + Z_v)/2.0

    # quality prediction
    alpha = -245.8909
    beta = 261.9373
    gamma1 = -239.8886 / 10000.0 
    gamma2 = 160.1664 / 10000.0 
    gamma3 = 64.2859 / 10000.0 

    # corner case of a black / white frame
    if np.abs(A) < 1e-3 or np.abs(B) < 1e-3 or np.abs(Z) < 1e-3:
        score = 0.0
    else:
        score = alpha + beta*(B**gamma1)*(A**gamma2)*(Z**gamma3)

    return score

def vis_matching(src_img, src_pts, dst_pts, M):
    """
    Visualize matching results
    """
    assert dst_pts.shape[0] == 2 and src_pts.shape[0]
    assert dst_pts.shape[1] == src_pts.shape[1]

    # draw the tracked pts
    img = src_img.copy()
    num_matches = src_pts.shape[1]
    for i in range(num_matches):
        cv2.circle(img, (src_pts[0, i], src_pts[1, i]), 3, (0,0,255), 1)
        cv2.line(img,(src_pts[0, i], src_pts[1, i]),
                    (dst_pts[0, i], dst_pts[1, i]),(0,255,0),1)

    return img


def distribute_keypoints(kps, img_size, grid_size, num_features):
    ''' 
    Sample interest points on 2D grid for better homography
    The results is a set of points evenly distributed over 2D
    '''
    # kp index and score
    kp_grids = np.zeros((grid_size, grid_size, num_features), dtype=np.int) - 1
    kp_scores = np.zeros((grid_size, grid_size, num_features), dtype=np.float32) - 1
    grid_width, grid_height = np.ceil(img_size / grid_size)

    # insert kp into grids
    for kp_ind, kp in enumerate(kps):
        x, y = kp.pt
        score = kp.response
        grid_x, grid_y = int(np.floor(x / grid_width)), int(np.floor(y / grid_height))

        # find the lowest scored kp in current grid
        slot = np.argmin(kp_scores[grid_x, grid_y, :])

        # replace slot with current kp if necessary
        if kp_scores[grid_x, grid_y, slot] < score:
            kp_grids[grid_x, grid_y, slot] = kp_ind
            kp_scores[grid_x, grid_y, slot] = score

    # valid list of kp
    kp_list = kp_grids.ravel().tolist()
    kp_list = [ind for ind in kp_list if ind != -1]

    return kp_list


def fit_homography(kp1, kp2, matches, img_size=None, src_img=None):
    ''' 
    Fit homography using matrix using forward/backward matching
    Homography is not going to be accurate anyway
    Tricks here to avoid bias over foreground
    '''
    M = np.zeros((3,3))
    matchesMask = []
    
    if img_size is None:
        # use all matches
        valid_matches = matches
    else:
        # distribute key points across 2D grids
        kp_list = distribute_keypoints(kp1, img_size, 32, 5)
        valid_matches = [m for m in matches if m.queryIdx in kp_list]

    # convert kp to pts
    src_pts = np.float32([ kp1[m.queryIdx].pt for m in valid_matches ]).reshape(-1,1,2)
    dst_pts = np.float32([ kp2[m.trainIdx].pt for m in valid_matches ]).reshape(-1,1,2)

    # call opencv for homography (with a conservative threshold)
    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.2)
    matchesMask = mask.ravel().tolist()

    # visualize matching results
    if src_img is not None:
        src_pts = src_pts.reshape(-1, 2).transpose()
        dst_pts = dst_pts.reshape(-1, 2).transpose()
        valid_ind = np.where(mask==1)[0]

        src_pts = src_pts[:, valid_ind]
        dst_pts = dst_pts[:, valid_ind]

        vis_img = vis_matching(src_img, src_pts, dst_pts, M)
        cv2.imshow('test', vis_img)
        cv2.waitKey(0)

    return M, matchesMask


def match_frames(prev_frame, curr_frame, params):
    """
    Match frames and return the homography
    ONLY if the motion estimation passed a series of tests
    """
    scale = params['fScale']
    num_feats = params['nFeat']
    min_num_kps = params['nMinKp']
    min_num_matches = params['nMinMatch']
    vh_thresh = params['fVhThresh']
    quality_thresh = params['fQualityThresh']
    quality_ratio = params['fQualityRatio']

    # init orb detector/matcher
    orb = cv2.ORB_create(num_feats)
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    
    # resize frames if necessary
    if np.abs(scale-1) > 1e-3:
        prev_frame = cv2.resize(prev_frame, (0,0), fx=scale, fy=scale)
        curr_frame = cv2.resize(curr_frame, (0,0), fx=scale, fy=scale)
    h, w, c = curr_frame.shape

    # check if prev_frame and curr_frame are blurry
    prev_score = image_quality(prev_frame)
    curr_score = image_quality(curr_frame)
    # print prev_score, curr_score

    if prev_score < quality_thresh or curr_score < quality_thresh:
        #print 'quality_thresh'
        return np.zeros((3,3))
    
    if abs(prev_score / (curr_score + 1e-6) - 1) > quality_ratio:
        #print 'quality_ratio'
        return np.zeros((3,3))

    # orb detect and descriptors
    kp1, des1 = orb.detectAndCompute(prev_frame, None)
    kp2, des2 = orb.detectAndCompute(curr_frame, None)

    # too few kps: textureless / blurring
    if len(kp1) < min_num_kps or len(kp2) < min_num_kps:
        print 'too few kps: textureless / blurring'
        return np.zeros((3,3))

    # forward / backward matching
    matches = bf.match(des1, des2)

    # not enough matching
    if len(matches) < min_num_matches:
        print 'not enough matching'
        return np.zeros((3,3))

    # fit homograhpy (with some tricks)
    M, mask = fit_homography( kp1, kp2, matches, np.asarray([w, h]), None)
    
    #  degenerated homography
    if len(mask) == 0:
        print 'degenerated homography'
        return np.zeros((3,3))

    # translational component
    vx, vy = np.abs(M[0,2]), np.abs(M[1,2])

    # moving too fast (violating homography assumption)
    if (vx > vh_thresh) or (vy > vh_thresh):
        print 'moving too fast'
        return np.zeros((3,3))

    return M


def process_vid_file(vid_file):
    """
    Process a single video file 
    Make sure opencv is compiled with ffmpeg
    """
    # dict that holds params for matching
    params = {}
    # scale for motion estimation
    params['fScale'] = 1.0
    # number of descriptors
    params['nFeat'] = 2000
    # minimal number of descriptors
    params['nMinKp'] = 50
    # minimal number of matched descriptors
    params['nMinMatch'] = 30
    # maximum offset in flow 
    params['fVhThresh'] = 32.0
    # params for image quality (motion blur)
    params['fQualityThresh'] = 5.0
    params['fQualityRatio'] = 0.2

    # get video stats
    video_name = os.path.basename(vid_file[:-4])
    video_name.replace(' ', '')
    output_frame_list = []
    output_homography = []
    frame_index = 0

    # open video file
    cap = cv2.VideoCapture(vid_file)
    if not cap.isOpened():
        print "Can not open video file: {:s}".format(vid_file)
        return frame_pairs
    
    # fetch the first frame
    ret, prev_frame = cap.read()
    if (not ret) or (prev_frame is None):
        print "Can not read video file: {:s}".format(vid_file)

    start = time.time()

    # loop over all frames
    while(True):

        # increase the counter
        frame_index += 1

        # read current frame
        ret, frame = cap.read()

        # end of file?
        if (not ret) or (prev_frame is None):
            break

        # match the frames
        M = match_frames(prev_frame, frame, params)

        # copy to previous frame
        prev_frame = frame.copy()

        # save the frame number and the homography
        if len(M) > 0:
            output_frame_list.append(frame_index)
            output_homography.append(M)
            if len(output_homography) == 1:
                output_homography.append(M)
    # timing
    end = time.time()
    print "Averge time per frame: {:2f} s. ".format(float(end-start)/frame_index)
    print "{:2d} valid frame pairs out of {:2d} total pairs".format(len(output_frame_list), frame_index-1)
   
    return output_frame_list, output_homography


if __name__ == '__main__':
    dataset = sys.argv[1]
    hom_path = '/media/meerahahn/General_Datasets/' + dataset + '/homography/'
    vid_path = '/media/meerahahn/General_Datasets/' + dataset + '/videos/'
    videos = [vid for vid in os.listdir(vid_path)]

    #video = '/media/meerahahn/General_Datasets/EDGTEA/videos/OP01-R02-TurkeySandwich.mp4'
    for vid in videos:
        print 'processing ' + vid
        video = vid_path + vid
        if ('.avi' in os.path.basename(video)) or ('.mp4' in os.path.basename(video)) \
                and os.path.exists(video):
            print "Processing video file: {:s}".format(video)
            output_frame_list, output_homography = process_vid_file(video)
            #print len(output_homography)
            np.save(hom_path + vid[:-3] + 'npy' , output_homography)
        else:
            print video + ' not found'

920fe85a35d0fe19c9e0f915eb4cdf64d1856ea1
import argparse
import sys
import os
import glob
import time
import cv2
import numpy as np
import shutil


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(
                        description='Extract homography from a list of videos')
    parser.add_argument('--video_list', dest='video_list_file',
                        help='a file that contains all videos (one video per line)',
                        default=None, type=str)
    parser.add_argument('--output_folder', dest='output_folder',
                        help='output folder to save selected video frames (root of the dataset)',
                        default=None, type=str)
 
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    return args

def image_quality(img):
    """
    Implementation of the paper
    "No-Reference Perceptual Quality Assessment of JPEG Compressed Images"
    Used as a quick filter of frames
    """
    # convert bgr image to gray -> float32
    score = 0.0
    if img is None:
        return score

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    x = gray.astype(np.float32)
    h, w = x.shape[0], x.shape[1]

    # horizontal
    d_h = x[:,1:] - x[:,:-1]
    w_bound = int(8*(np.floor(w/8.0)-1)) + 1
    B_h = np.mean(np.abs(d_h[:,7:w_bound:8]))
    A_h = (8.0 * np.mean(np.abs(d_h)) - B_h) / 7.0
    sig_h = np.sign(d_h)
    left_sig, right_sig = sig_h[:,:-2], sig_h[:,1:-1]
    Z_h = np.mean((left_sig * right_sig)<0)

    # vertical
    d_v = x[1:, :] - x[:-1, :]
    h_bound = int(8*(np.floor(h/8.0)-1)) + 1
    B_v = np.mean(np.abs(d_v[7:h_bound:8, :]))
    A_v = (8.0 * np.mean(np.abs(d_v)) - B_v) / 7.0
    sig_v = np.sign(d_v)
    up_sig, down_sig = sig_v[:-2, :], sig_v[1:-1, :]
    Z_v = np.mean((up_sig * down_sig)<0)

    # combine the weights
    B = (B_h + B_v)/2.0
    A = (A_h + A_v)/2.0
    Z = (Z_h + Z_v)/2.0

    # quality prediction
    alpha = -245.8909
    beta = 261.9373
    gamma1 = -239.8886 / 10000.0 
    gamma2 = 160.1664 / 10000.0 
    gamma3 = 64.2859 / 10000.0 

    # corner case of a black / white frame
    if np.abs(A) < 1e-3 or np.abs(B) < 1e-3 or np.abs(Z) < 1e-3:
        score = 0.0
    else:
        score = alpha + beta*(B**gamma1)*(A**gamma2)*(Z**gamma3)

    return score

def vis_matching(src_img, src_pts, dst_pts, M):
    """
    Visualize matching results
    """
    assert dst_pts.shape[0] == 2 and src_pts.shape[0]
    assert dst_pts.shape[1] == src_pts.shape[1]

    # draw the tracked pts
    img = src_img.copy()
    num_matches = src_pts.shape[1]
    for i in range(num_matches):
        cv2.circle(img, (src_pts[0, i], src_pts[1, i]), 3, (0,0,255), 1)
        cv2.line(img,(src_pts[0, i], src_pts[1, i]),
                    (dst_pts[0, i], dst_pts[1, i]),(0,255,0),1)

    return img


def distribute_keypoints(kps, img_size, grid_size, num_features):
    ''' 
    Sample interest points on 2D grid for better homography
    The results is a set of points evenly distributed over 2D
    '''
    # kp index and score
    kp_grids = np.zeros((grid_size, grid_size, num_features), dtype=np.int) - 1
    kp_scores = np.zeros((grid_size, grid_size, num_features), dtype=np.float32) - 1
    grid_width, grid_height = np.ceil(img_size / grid_size)

    # insert kp into grids
    for kp_ind, kp in enumerate(kps):
        x, y = kp.pt
        score = kp.response
        grid_x, grid_y = int(np.floor(x / grid_width)), int(np.floor(y / grid_height))

        # find the lowest scored kp in current grid
        slot = np.argmin(kp_scores[grid_x, grid_y, :])

        # replace slot with current kp if necessary
        if kp_scores[grid_x, grid_y, slot] < score:
            kp_grids[grid_x, grid_y, slot] = kp_ind
            kp_scores[grid_x, grid_y, slot] = score

    # valid list of kp
    kp_list = kp_grids.ravel().tolist()
    kp_list = [ind for ind in kp_list if ind != -1]

    return kp_list


def fit_homography(kp1, kp2, matches, img_size=None, src_img=None):
    ''' 
    Fit homography using matrix using forward/backward matching
    Homography is not going to be accurate anyway
    Tricks here to avoid bias over foreground
    '''
    M = np.zeros((3,3))
    matchesMask = []
    
    if img_size is None:
        # use all matches
        valid_matches = matches
    else:
        # distribute key points across 2D grids
        kp_list = distribute_keypoints(kp1, img_size, 32, 5)
        valid_matches = [m for m in matches if m.queryIdx in kp_list]

    # convert kp to pts
    src_pts = np.float32([ kp1[m.queryIdx].pt for m in valid_matches ]).reshape(-1,1,2)
    dst_pts = np.float32([ kp2[m.trainIdx].pt for m in valid_matches ]).reshape(-1,1,2)

    # call opencv for homography (with a conservative threshold)
    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.2)
    matchesMask = mask.ravel().tolist()

    # visualize matching results
    if src_img is not None:
        src_pts = src_pts.reshape(-1, 2).transpose()
        dst_pts = dst_pts.reshape(-1, 2).transpose()
        valid_ind = np.where(mask==1)[0]

        src_pts = src_pts[:, valid_ind]
        dst_pts = dst_pts[:, valid_ind]

        vis_img = vis_matching(src_img, src_pts, dst_pts, M)
        cv2.imshow('test', vis_img)
        cv2.waitKey(0)

    return M, matchesMask


def match_frames(prev_frame, curr_frame, params):
    """
    Match frames and return the homography
    ONLY if the motion estimation passed a series of tests
    """
    scale = params['fScale']
    num_feats = params['nFeat']
    min_num_kps = params['nMinKp']
    min_num_matches = params['nMinMatch']
    vh_thresh = params['fVhThresh']
    quality_thresh = params['fQualityThresh']
    quality_ratio = params['fQualityRatio']

    # init orb detector/matcher
    orb = cv2.ORB_create(num_feats)
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    
    # resize frames if necessary
    if np.abs(scale-1) > 1e-3:
        prev_frame = cv2.resize(prev_frame, (0,0), fx=scale, fy=scale)
        curr_frame = cv2.resize(curr_frame, (0,0), fx=scale, fy=scale)
    h, w, c = curr_frame.shape

    # check if prev_frame and curr_frame are blurry
    prev_score = image_quality(prev_frame)
    curr_score = image_quality(curr_frame)
    # print prev_score, curr_score

    if prev_score < quality_thresh or curr_score < quality_thresh:
        #print 'quality_thresh'
        return np.zeros((3,3))
    
    if abs(prev_score / (curr_score + 1e-6) - 1) > quality_ratio:
        #print 'quality_ratio'
        return np.zeros((3,3))

    # orb detect and descriptors
    kp1, des1 = orb.detectAndCompute(prev_frame, None)
    kp2, des2 = orb.detectAndCompute(curr_frame, None)

    # too few kps: textureless / blurring
    if len(kp1) < min_num_kps or len(kp2) < min_num_kps:
        print 'too few kps: textureless / blurring'
        return np.zeros((3,3))

    # forward / backward matching
    matches = bf.match(des1, des2)

    # not enough matching
    if len(matches) < min_num_matches:
        print 'not enough matching'
        return np.zeros((3,3))

    # fit homograhpy (with some tricks)
    M, mask = fit_homography( kp1, kp2, matches, np.asarray([w, h]), None)
    
    #  degenerated homography
    if len(mask) == 0:
        print 'degenerated homography'
        return np.zeros((3,3))

    # translational component
    vx, vy = np.abs(M[0,2]), np.abs(M[1,2])

    # moving too fast (violating homography assumption)
    if (vx > vh_thresh) or (vy > vh_thresh):
        print 'moving too fast'
        return np.zeros((3,3))

    return M


def process_vid_file(vid_file):
    """
    Process a single video file 
    Make sure opencv is compiled with ffmpeg
    """
    # dict that holds params for matching
    params = {}
    # scale for motion estimation
    params['fScale'] = 1.0
    # number of descriptors
    params['nFeat'] = 2000
    # minimal number of descriptors
    params['nMinKp'] = 50
    # minimal number of matched descriptors
    params['nMinMatch'] = 30
    # maximum offset in flow 
    params['fVhThresh'] = 32.0
    # params for image quality (motion blur)
    params['fQualityThresh'] = 5.0
    params['fQualityRatio'] = 0.2

    # get video stats
    video_name = os.path.basename(vid_file[:-4])
    video_name.replace(' ', '')
    output_frame_list = []
    output_homography = []
    frame_index = 0

    # open video file
    cap = cv2.VideoCapture(vid_file)
    if not cap.isOpened():
        print "Can not open video file: {:s}".format(vid_file)
        return frame_pairs
    
    # fetch the first frame
    ret, prev_frame = cap.read()
    if (not ret) or (prev_frame is None):
        print "Can not read video file: {:s}".format(vid_file)

    start = time.time()

    # loop over all frames
    while(True):

        # increase the counter
        frame_index += 1

        # read current frame
        ret, frame = cap.read()

        # end of file?
        if (not ret) or (prev_frame is None):
            break

        # match the frames
        M = match_frames(prev_frame, frame, params)

        # copy to previous frame
        prev_frame = frame.copy()

        # save the frame number and the homography
        if len(M) > 0:
            output_frame_list.append(frame_index)
            output_homography.append(M)
            if len(output_homography) == 1:
                output_homography.append(M)
    # timing
    end = time.time()
    print "Averge time per frame: {:2f} s. ".format(float(end-start)/frame_index)
    print "{:2d} valid frame pairs out of {:2d} total pairs".format(len(output_frame_list), frame_index-1)
   
    return output_frame_list, output_homography


if __name__ == '__main__':
    dataset = sys.argv[1]
    hom_path = '/media/meerahahn/General_Datasets/' + dataset + '/homography/'
    vid_path = '/media/meerahahn/General_Datasets/' + dataset + '/videos/'
    videos = [vid for vid in os.listdir(vid_path)]

    #video = '/media/meerahahn/General_Datasets/EDGTEA/videos/OP01-R02-TurkeySandwich.mp4'
    for vid in videos:
        print 'processing ' + vid
        video = vid_path + vid
        if ('.avi' in os.path.basename(video)) or ('.mp4' in os.path.basename(video)) \
                and os.path.exists(video):
            print "Processing video file: {:s}".format(video)
            output_frame_list, output_homography = process_vid_file(video)
            #print len(output_homography)
            np.save(hom_path + vid[:-3] + 'npy' , output_homography)
        else:
            print video + ' not found'
