fuse_upsamp_gen_flow.py

import os

import numpy as np
import os
import torch
import torch.nn as nn
import torch.utils.data as data 
import cv2
import argparse

from fuse_upsampled_data.flow_generator import UpsampFlowFixNGenerator, GTFlowFixNGenerator

def events_to_save_format(events, t0):
    t = (events[:, 0]).astype(np.float32)
    if t0:
        t -= t0
    x = events[:, 1].astype(np.int16)
    y = events[:, 2].astype(np.int16)
    
    if events[:, 3].min() == -1:
        p = ((events[:, 3] + 1) / 2).astype(np.bool_) # go from -1/+1 to 0,1, before it is int16
    else:
        p = events[:, 3].astype(np.bool_) 
    return t, x, y, p


def save_fused_events_frames_flows(path_to_sequence, seq_events, seq_frames, seq_timestamps, seq_flows):
    path_to_events = os.path.join(path_to_sequence, 'events')
    path_to_frames = os.path.join(path_to_sequence, 'frames')
    path_to_flows = os.path.join(path_to_sequence, 'flow')
    path_to_save_timestamps = os.path.join(path_to_frames, 'timestamps.txt')
    
    
    if not os.path.exists(path_to_flows):
        os.makedirs(path_to_events)
        os.makedirs(path_to_frames)
        os.makedirs(path_to_flows)
    
    f_timestamps = open(path_to_save_timestamps, 'w+')

    # t0 = 0
    # if is_t0:
    t0 = seq_timestamps[0]
    seq_timestamps = np.array(seq_timestamps) - t0

    NE = 0
    for i in range(len(seq_frames)):
        
        if i < len(seq_events):
            t,x,y,p = events_to_save_format(seq_events[i], t0)
            # print('events {}:'.format(i), t[0],t[-1])
            np.savez_compressed(os.path.join(path_to_events, 'events_{:010d}.npz'.format(i)), t=t, x=x, y=y, p=p)
            
            np.savez_compressed(os.path.join(path_to_flows, 'flow_{:010d}.npz'.format(i)), flow01=seq_flows[i][:2,:,:], flow10=seq_flows[i][2:,:,:])
            NE += len(t)
        cv2.imwrite(os.path.join(path_to_frames, 'frame_{:010d}.png'.format(i)), seq_frames[i])
        f_timestamps.write(str(i)+' '+str(seq_timestamps[i])+'\n')
        
    f_timestamps.close()
    print('Avg Event number: {:.0f}, Start events: {:.5f}, End events: {:.5f}'.format(NE/len(seq_frames), seq_events[0][0][0], seq_events[-1][-1][0]))
    print('Start frame: {:.5f}, End frame: {:.5f} '.format(seq_timestamps[0]+t0, seq_timestamps[-1]+t0))


if __name__ == "__main__":
    '''
    Fuse events with fixed event count, and generate corresponding flow
    Two options:
    1. is_gt_flow == True and motion config files are known (specified in data_txt_file file) 
        Generate GT flow based on motion using VideoRenderer.generate_gt_flow(t0, t1)
    2. is_gt_flow == False or motions are not known
        Generate estimated flow based on FlowNet
    '''
    parser = argparse.ArgumentParser(description='Generate intensity video sequences')
    
    parser.add_argument('--data_txt_file', type=str, required=True, help='path to txt file for loading data, generated by generate_txt.py')
    parser.add_argument('--path_to_read_data', type=str, required=True, help='path to HFR data, including events/frames')
    parser.add_argument('--path_to_save_data', type=str, required=True, help='path to save data')
    parser.add_argument('--len_sequence', type=int, default=10)
    parser.add_argument('--max_n_per_seq', type=int, default=-1, help='max number of splited sequences per original sequence')
    parser.add_argument('--num_events', type=int, default=15000)
    parser.add_argument('--image_dim', nargs=2, type=int, default=[180, 240], help="[height, width]")
    # parser.add_argument('--is_t0', dest='is_t0', action='store_true', default=False)
    parser.add_argument('--is_gt_flow', dest='is_gt_flow', action='store_true', default=False)
    #---- Replace root path in motion configs, in case the root path changes ------
    parser.add_argument('--old_root_path', type=str, help='Old root dir to read frames', \
        default=None)
    parser.add_argument('--new_root_path', type=str, help='New root dir to read frames', \
        default=None)
    #-------------------------------
    parser.add_argument('--start_id', type=int, default=0)
    parser.add_argument('--end_id', type=int, default=-1)
    
    args = parser.parse_args()
    
    if args.is_gt_flow:
        flow_generator = GTFlowFixNGenerator(args.data_txt_file, args.path_to_read_data, cfgs=args, device='cuda:0', \
             replace_path_pair=[args.old_root_path, args.new_root_path] if args.new_root_path else None)
    else:
        flow_generator = UpsampFlowFixNGenerator(args.data_txt_file, args.path_to_read_data, cfgs=args, device='cuda:0')
    

    if not os.path.exists(args.path_to_save_data):
        os.makedirs(args.path_to_save_data)
    
    
    print('Number of fused sequences: ', len(flow_generator))
    for batch_idx, cur_data in enumerate(flow_generator):
        seq_events, seq_frames, seq_timestamps, seq_flows = cur_data

        path_to_sequence = os.path.join(args.path_to_save_data, 'sequence_{:010d}'.format(batch_idx))
        if not os.path.exists(path_to_sequence):
            os.makedirs(path_to_sequence)
        print('Processing {} ...'.format(path_to_sequence))
        

        save_fused_events_frames_flows(path_to_sequence, seq_events, seq_frames, seq_timestamps, seq_flows) #, args.is_t0