Dronet-Pytorch/move_data.py at master · peasant98/Dronet-Pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import numpy as np

import os
import pandas as pd
from enum import Enum

# data module for dronet.

def rand_split(data_path_folder):
    # need to move the data into correct spot first
    csv_path = os.path.join(data_path_folder, 'interpolated.csv')
    # results = os.listdir(data_path_folder)
    interpolated_df = pd.read_csv(csv_path, usecols=[1,2,3,4,5,6,7,8])
    # get the frames that are center camera
    correct_interpolated_df = interpolated_df[interpolated_df['frame_id']=='center_camera']
    correct_interpolated_df['filename'] = correct_interpolated_df['filename'].apply(lambda x: 'images'+x[6:])

    img_list = os.listdir(os.path.join(data_path_folder, 'images'))
    real_indices = ((correct_interpolated_df.index.values - 2)/3).astype(int)
    correct_interpolated_df = correct_interpolated_df.assign(norm_index=real_indices)
    # randomly permute values for split of train, test, and validation
    # split is 70 10 20
    res = correct_interpolated_df.iloc[np.random.permutation(len(correct_interpolated_df))]
    selection = res['norm_index'].values[0]

    length = len(res)
    # get splits
    first_split_idx = int(length * 0.7)
    remaining = length - first_split_idx
    second_split_idx = int(first_split_idx + (remaining*0.333))
    train = res.iloc[:first_split_idx]
    valid = res.iloc[first_split_idx:second_split_idx]
    test = res.iloc[second_split_idx:]
    # return split dataframes
    train.to_csv('training_steer.csv')
    valid.to_csv('valid_steer.csv')
    test.to_csv('test_steer.csv')
    return train, valid, test


def move_img(df, data_group, udacity_data_path_folder='out', main_data_file_path='data/collision/collision_dataset'):
    '''
    moves the image and adds to steering.txt in order.
    '''
    udacity_name = 'UDACITY_{}'.format(data_group.upper())
    arr = np.zeros(len(df))
    for ind,filename in enumerate(df['filename'].values):
        # get initial filename
        # name by ind
        # example filename is images/341401.jpg
        # rename to images/{ind}.jpg
        new_filename = os.path.join('images', '{}.jpg'.format(ind))
        fname = os.path.join(udacity_data_path_folder, filename)
        print(os.path.join(main_data_file_path, data_group, udacity_name, new_filename))
        # move to main data folder
        os.rename(fname, os.path.join(main_data_file_path, data_group, udacity_name, new_filename))
        # add to file
        angle = df['angle'].values[ind]
        arr[ind] = angle
    steering_fname = 'steering_{}.txt'.format(data_group)
    np.savetxt(steering_fname, arr)
    os.rename(steering_fname, os.path.join(main_data_file_path, data_group, udacity_name, steering_fname))

if __name__ == '__main__':
    train_df, valid_df, test_df = rand_split('out')
    print(train_df.head(), valid_df.head(), test_df.head())
    move_img(train_df, 'training')
    move_img(valid_df, 'validation')
    move_img(test_df, 'testing')