Skip to content
GitLab
Switch branch/tag
Find file BlameHistoryPermalink
data.py 8.46 KiB
1 
"""
2 
Class for managing our data.
3 
"""
4 
import csv
5 
import numpy as np
6 
import random
7 
import glob
8 
import os.path
9 
import pandas as pd
10 
import sys
11 
import operator
12 
from processor import process_image
13 
from keras.utils import np_utils
14
15 
class DataSet():
16
17 
    def __init__(self, seq_length=40, class_limit=None, image_shape=(224, 224, 3)):
18 
        """Constructor.
19 
        seq_length = (int) the number of frames to consider
20 
        class_limit = (int) number of classes to limit the data to.
21 
            None = no limit.
22 
        """
23 
        self.seq_length = seq_length
24 
        self.class_limit = class_limit
25 
        self.sequence_path = os.path.join('data', 'sequences')
26 
        self.max_frames = 300  # max number of frames a video can have for us to use it
27
28 
        # Get the data.
29 
        self.data = self.get_data()
30
31 
        # Get the classes.
32 
        self.classes = self.get_classes()
33
34 
        # Now do some minor data cleaning.
35 
        self.data = self.clean_data()
36
37 
        self.image_shape = image_shape
38
39 
    @staticmethod
40 
    def get_data():
41 
        """Load our data from file."""
42 
        with open(os.path.join('data', 'data_file.csv'), 'r') as fin:
43 
            reader = csv.reader(fin)
44 
            data = list(reader)
45
46 
        return data
47
48 
    def clean_data(self):
49 
        """Limit samples to greater than the sequence length and fewer
50 
        than N frames. Also limit it to classes we want to use."""
51 
        data_clean = []
52 
        for item in self.data:
53 
            if int(item[3]) >= self.seq_length and int(item[3]) <= self.max_frames \
54 
                    and item[1] in self.classes:
55 
                data_clean.append(item)
56
57 
        return data_clean
58
59 
    def get_classes(self):
60 
        """Extract the classes from our data. If we want to limit them,
61 
        only return the classes we need."""
62 
        classes = []
63 
        for item in self.data:
64 
            if item[1] not in classes:
65 
                classes.append(item[1])
66
67 
        # Sort them.
68 
        classes = sorted(classes)
69
70 
        # Return.
7172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
if self.class_limit is not None: return classes[:self.class_limit] else: return classes def get_class_one_hot(self, class_str): """Given a class as a string, return its number in the classes list. This lets us encode and one-hot it for training.""" # Encode it first. label_encoded = self.classes.index(class_str) # Now one-hot it. label_hot = np_utils.to_categorical(label_encoded, len(self.classes)) label_hot = label_hot[0] # just get a single row return label_hot def split_train_test(self): """Split the data into train and test groups.""" train = [] test = [] for item in self.data: if item[0] == 'train': train.append(item) else: test.append(item) return train, test def get_all_sequences_in_memory(self, train_test, data_type, concat=False): """ This is a mirror of our generator, but attempts to load everything into memory so we can train way faster. """ # Get the right dataset. train, test = self.split_train_test() data = train if train_test == 'train' else test print("Loading %d samples into memory for %sing." % (len(data), train_test)) X, y = [], [] for row in data: if data_type == 'images': frames = self.get_frames_for_sample(row) frames = self.rescale_list(frames, self.seq_length) # Build the image sequence sequence = self.build_image_sequence(frames) else: sequence = self.get_extracted_sequence(data_type, row) if sequence is None: print("Can't find sequence. Did you generate them?") raise if concat: # We want to pass the sequence back as a single array. This # is used to pass into a CNN or MLP, rather than an RNN. sequence = np.concatenate(sequence).ravel() X.append(sequence) y.append(self.get_class_one_hot(row[1])) return np.array(X), np.array(y) def frame_generator(self, batch_size, train_test, data_type, concat=False): """Return a generator that we can use to train on. There are a couple different things we can return:
141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
data_type: 'features', 'images' """ # Get the right dataset for the generator. train, test = self.split_train_test() data = train if train_test == 'train' else test print("Creating %s generator with %d samples." % (train_test, len(data))) while 1: X, y = [], [] # Generate batch_size samples. for _ in range(batch_size): # Reset to be safe. sequence = None # Get a random sample. sample = random.choice(data) # Check to see if we've already saved this sequence. if data_type is "images": # Get and resample frames. frames = self.get_frames_for_sample(sample) frames = self.rescale_list(frames, self.seq_length) # Build the image sequence sequence = self.build_image_sequence(frames) else: # Get the sequence from disk. sequence = self.get_extracted_sequence(data_type, sample) if sequence is None: print("Can't find sequence. Did you generate them?") sys.exit() # TODO this should raise if concat: # We want to pass the sequence back as a single array. This # is used to pass into an MLP rather than an RNN. sequence = np.concatenate(sequence).ravel() X.append(sequence) y.append(self.get_class_one_hot(sample[1])) yield np.array(X), np.array(y) def build_image_sequence(self, frames): """Given a set of frames (filenames), build our sequence.""" return [process_image(x, self.image_shape) for x in frames] def get_extracted_sequence(self, data_type, sample): """Get the saved extracted features.""" filename = sample[2] path = self.sequence_path + filename + '-' + str(self.seq_length) + \ '-' + data_type + '.txt' if os.path.isfile(path): # Use a dataframe/read_csv for speed increase over numpy. features = pd.read_csv(path, sep=" ", header=None) return features.values else: return None @staticmethod def get_frames_for_sample(sample): """Given a sample row from the data file, get all the corresponding frame filenames.""" path = os.path.join('data', sample[0], sample[1]) filename = sample[2] images = sorted(glob.glob(os.path.join(path, filename + '*jpg'))) return images
211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252
@staticmethod def get_filename_from_image(filename): parts = filename.split(os.path.sep) return parts[-1].replace('.jpg', '') @staticmethod def rescale_list(input_list, size): """Given a list and a size, return a rescaled/samples list. For example, if we want a list of size 5 and we have a list of size 25, return a new list of size five which is every 5th element of the origina list.""" assert len(input_list) >= size # Get the number to skip between iterations. skip = len(input_list) // size # Build our new output. output = [input_list[i] for i in range(0, len(input_list), skip)] # Cut off the last one if needed. return output[:size] @staticmethod def print_class_from_prediction(predictions, nb_to_return=5): """Given a prediction, print the top classes.""" # Get the prediction for each label. label_predictions = {} for i, label in enumerate(data.classes): label_predictions[label] = predictions[i] # Now sort them. sorted_lps = sorted( label_predictions.items(), key=operator.itemgetter(1), reverse=True ) # And return the top N. for i, class_prediction in enumerate(sorted_lps): if i > nb_to_return - 1 or class_prediction[1] == 0.0: break print("%s: %.2f" % (class_prediction[0], class_prediction[1]))