dsafdpi/create_data.py at master · muhammadfaiz12/dsafdpi · 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
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
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

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 = ",trainPath)
    videoPaths = glob(trainPath+'/*')
    print("videPath = ",videoPaths)
    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 = "train/"+genre+"_train_"+model_name
    dump_pkl(genreFeatures, outPath)

def gather_testing_data(genre, model_name=default_model_name):
    """Driver function to collect frame features for a genre"""

    trainPath = os.path.join(video_resource,'test',genre)
    print("trainpath = ",trainPath)
    videoPaths = glob(trainPath+'/*')
    print("videPath = ",videoPaths)
    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 = "test/"+genre+"_test_"+model_name
    dump_pkl(genreFeatures, outPath)


if __name__=="__main__":
    from sys import argv
    extractor = argv[1]
    genres = ['action','horror','romance']#
    for genre in genres:
        gather_training_data(genre, extractor)
        gather_testing_data(genre, extractor)