forked from maximus009/MovieScope
-
Notifications
You must be signed in to change notification settings - Fork 0
/
create_data.py
37 lines (31 loc) · 1.07 KB
/
create_data.py
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
import os, sys
from glob import glob
import cv2
import numpy as np
from config.global_parameters import default_model_name
from config.resources import video_resource
from model_utils import get_features_batch
from utils import dump_pkl
from video import get_frames
def gather_training_data(genre, model_name=default_model_name):
"""Driver function to collect frame features for a genre"""
trainPath = os.path.join(video_resource,'train',genre)
print trainPath
videoPaths = glob(trainPath+'/*')
genreFeatures = []
for videoPath in videoPaths:
print videoPath,":",
frames =list(get_frames(videoPath, time_step=1000))
print len(frames),
if len(frames)==0:
print "corrupt."
continue
videoFeatures = get_features_batch(frames, model_name)
print videoFeatures.shape
genreFeatures.append(videoFeatures)
outPath = genre+"_ultimate_"+model_name
dump_pkl(genreFeatures, outPath)
if __name__=="__main__":
from sys import argv
genre = argv[1]
gather_training_data(genre)