extract_features.py

import pandas as pd
import numpy as np
import pickle
import os

# from feature_models.vggish import vggish_input, vggish_params, vggish_slim
import features.config as fconf
from features.audio import extract_VGGish_features
from features.image import extract_resnet152_features
from features.text import fit_tokenizer, extract_GLoVe_features
from features.emotion import extract_emotions

DATASETS = ["training_set", "development_set", "testing_set"]
CAPTION_FILES = {
    "training_set": "train_text_descriptions.csv",
    "development_set": "dev_text_descriptions.csv",
    "testing_set": "test_text_descriptions.csv"
}
FEATURE_MODEL_DIR = "features/"


def get_boolean_input(prompt):
    while True:
        ans = input(prompt).strip().lower()
        if ans == "y":
            return True
        elif ans == "n":
            return False
        print("Invalid response.")


def extract_features(
        datasets=DATASETS,
        caption_files=CAPTION_FILES,
        frame_subdir="Frames",
        feature_model_dir=FEATURE_MODEL_DIR,
        audio_subdir="Audio"):

    # extract_all = get_boolean_input(
    #     "Extract all features? (y/n) Otherwise extract on per-feature basis.")

    # if extract_all or get_boolean_input("Extract emotion features? (y/n)"):
    #     for dataset in datasets:
    #         feature_dir = fconf.set_dataset(dataset, fconf.EMOTION_FEATURE_DIR)
    #         print(f"Extracting emotion features for {dataset}...")
    #         _ = extract_emotions(
    #             frame_path=f"{dataset}/{frame_subdir}", feature_dir=feature_dir)
    #     print("Finished emotion feature extraction.\n")

    # if extract_all or get_boolean_input("Extract text (GLoVe) features? (y/n)"):
    caption_paths = [
        f"{dataset}/{caption_files[dataset]}" for dataset in datasets]
    all_caption_data = pd.concat(
        [pd.read_csv(path) for path in caption_paths])
    all_video_ids = all_caption_data["video_id"]
    all_captions = all_caption_data["description"]

    caption_count = all_video_ids.groupby(all_video_ids).count()
    print(
        f"\tFound {len(all_captions)} captions for {len(all_video_ids.unique())} videos")
    print("\tAvg num of captions per video:", np.mean(caption_count))
    print("\tMax num of captions per video:", np.max(caption_count))
    print("\tMin num of captions per video:", np.min(caption_count))

    sequence_lengths = [len(caption.split()) for caption in all_captions]
    max_sequence_length = np.max(sequence_lengths)

    print("Max sequence length:", max_sequence_length)
    print("Min sequence length:", np.min(sequence_lengths))
    print("Avg sequence length:", np.mean(sequence_lengths))

    print("Fitting tokenizer...")
    tokenizer = fit_tokenizer(all_captions)

    for dataset, caption_path in zip(datasets, caption_paths):
        print(f"Extracting GLoVe features from {caption_path}...")
        caption_data = pd.read_csv(caption_path)
        video_ids = caption_data["video_id"]
        captions = caption_data["description"]
        feature_dir = fconf.set_dataset(dataset, fconf.GLOVE_FEATURE_DIR)
        extract_GLoVe_features(
            video_ids=video_ids,
            captions=captions,
            tokenizer=tokenizer,
            max_sequence_length=max_sequence_length,
            feature_dir=feature_dir)

    tokenizer_pickle_path = f"{feature_model_dir}/caption_tokenizer.pickle"
    print(f"Saving tokenizer at {tokenizer_pickle_path}")
    with open(tokenizer_pickle_path, "wb") as f:
        pickle.dump(tokenizer, f, protocol=pickle.HIGHEST_PROTOCOL)
    print("Finished GLoVe feature extraction.\n")

    # if extract_all or get_boolean_input("Extract image (ResNet152) features? (y/n)"):
    #     for dataset in datasets:
    #         frame_dir = f"{dataset}/{frame_subdir}"
    #         feature_dir = fconf.set_dataset(
    #             dataset, fconf.RESNET152_FEATURE_DIR)
    #         print(f"Extracting ResNet152 features for {dataset}...")
    #         extract_resnet152_features(
    #             image_dir=frame_dir, features_dir=feature_dir)
    #     print("Finished ResNet152 feature extraction.\n")

    # if extract_all or get_boolean_input("Extract audio (VGGish) features? (y/n)"):
    #     for dataset in datasets:
    #         audio_dir = f"{dataset}/{audio_subdir}"
    #         video_ids = [f.split(".")[0]
    #                      for f in os.listdir(audio_dir) if f[-4:] == ".wav"]
    #         feature_dir = fconf.set_dataset(dataset, fconf.VGGISH_FEATURE_DIR)
    #         print(f"Extracting VGGish features for {dataset}...")
    #         extract_VGGish_features(
    #             video_ids,
    #             vggish_slim=vggish_slim,
    #             vggish_params=vggish_params,
    #             vggish_input=vggish_input,
    #             model_ckpt_path=f"{feature_model_dir}/vggish/vggish_model.ckpt",
    #             embedding_dir=feature_dir,
    #             audio_dir=audio_dir)
    #     print("Finished VGGish feature extraction.\n")


if __name__ == "__main__":
    extract_features()