run_s2vt.py

'''
A short script to train and evaluate s2vt model on Flickr30k and
MSR-VTT datasets.

Usage:
            python run_s2vt.py --dataset [MSR-VTT|Flickr30k] --train
            python run_s2vt.py --dataset [MSR-VTT|Flickr30k] --test --checkpoint model_num 
'''

from cfg import *
import os

import pandas as pd
import numpy as np
import pickle as pkl
import json
from tqdm import tqdm
import argparse

from pandas.io.json import json_normalize
from cocoeval import COCOScorer, suppress_stdout_stderr
import sys

from s2vt_model import *

cfg = None


def get_flickr30k_data(cfg):
    #using the provided splits
    train_split = set(map(lambda x: x.split(".")[0], open(cfg.train_file).read().splitlines()))
    val_split = set(map(lambda x: x.split(".")[0], open(cfg.val_file).read().splitlines()))
    test_split = set(map(lambda x: x.split(".")[0], open(cfg.test_file).read().splitlines()))
    
    data = [{"video_id": item.split(".")[0], "sentence_id": item.split("#")[1].split("\t")[0], "caption":item.split("\t")[1]}
            for item in open(cfg.annotations_path).read().splitlines()]
    
    sentences = json_normalize(data)
    sentences['video_path'] = sentences['video_id'].map(lambda x: os.path.join(cfg.path_to_descriptors, x + cfg.descriptor_suffix + ".npy"))
    
    train_imgs = sentences.loc[sentences["video_id"].isin(train_split)]
    train_imgs.reset_index()
    
    val_imgs = sentences.loc[sentences["video_id"].isin(val_split)]
    val_imgs.reset_index()
    
    test_imgs = sentences.loc[sentences["video_id"].isin(test_split)]
    test_imgs.reset_index()
    
    return train_imgs, val_imgs, test_imgs

def get_msr_vtt_data(cfg):
    #trainval data
    with open(cfg.trainval_annotations) as data_file:    
        data = json.load(data_file)
    
    sentences = json_normalize(data['sentences'])
    videos = json_normalize(data['videos'])
    train_vids = sentences.loc[sentences["video_id"].isin(videos[videos['split'] == "train"]["video_id"])]
    val_vids = sentences.loc[sentences["video_id"].isin(videos[videos['split'] == "validate"]["video_id"])]
    train_vids['video_path'] = train_vids['video_id'].map(lambda x: os.path.join(cfg.path_to_trainval_descriptors, x + "_incp_v3.npy"))  
    val_vids['video_path'] = val_vids['video_id'].map(lambda x: os.path.join(cfg.path_to_trainval_descriptors, x + "_incp_v3.npy"))
    
    #test data
    with open(cfg.test_annotations) as data_file:    
        data = json.load(data_file)
    sentences = json_normalize(data['sentences'])
    videos = json_normalize(data['videos'])
    test_vids = sentences.loc[sentences["video_id"].isin(videos[videos['split'] == "test"]["video_id"])]
    test_vids['video_path'] = test_vids['video_id'].map(lambda x: os.path.join(cfg.path_to_test_descriptors, x + "_incp_v3.npy"))
    
    return train_vids, val_vids, test_vids
    

def preProBuildWordVocab(sentence_iterator, word_count_threshold=5): # borrowed this function from NeuralTalk
    print 'preprocessing word counts and creating vocab based on word count threshold %d' % (word_count_threshold, )
    word_counts = {}
    nsents = 0
    for sent in sentence_iterator:
        nsents += 1
        for w in sent.lower().split(' '):
            word_counts[w] = word_counts.get(w, 0) + 1

    vocab = [w for w in word_counts if word_counts[w] >= word_count_threshold]
    print 'filtered words from %d to %d' % (len(word_counts), len(vocab))

    ixtoword = {}
    ixtoword[0] = '.'  # period at the end of the sentence. make first dimension be end token
    wordtoix = {}
    wordtoix['#START#'] = 0 # make first vector be the start token
    ix = 1
    for w in vocab:
        wordtoix[w] = ix
        ixtoword[ix] = w
        ix += 1

    word_counts['.'] = nsents
    bias_init_vector = np.array([1.0*word_counts[ixtoword[i]] for i in ixtoword])
    bias_init_vector /= np.sum(bias_init_vector) # normalize to frequencies
    bias_init_vector = np.log(bias_init_vector)
    bias_init_vector -= np.max(bias_init_vector) # shift to nice numeric range
    return wordtoix, ixtoword, bias_init_vector

def output_progress(current, total, loss):
    bar_length = 20
    progress = current/float(total)
    sys.stdout.write('\r')
    sys.stdout.write(("[%-" + str(bar_length) + "s] %d/%d") % ('='* int(bar_length * progress) + ">", current, total) + ", avg_loss=" + str(loss))
    sys.stdout.flush()

#populate feature dictionary
#unroll features for LSTM encoding
feature_dict = {} 
def load_flickr30k_features(vid):
    if vid in feature_dict:
        return feature_dict[vid]
    else:
        temp_array = np.load(vid)
        temp_array[1::2][:] = temp_array[1::2][:, ::-1][:]
        if cfg.use_hard_cache:
            feature_dict[vid] = temp_array.reshape(cfg.n_frame_step, -1)
            return feature_dict[vid]
        else:
            return temp_array.reshape(cfg.n_frame_step, -1)

def load_msr_vtt_features(vid):
    return np.load(vid)
            

def get_validation_loss(sess, current_val_data, wordtoix, tf_loss, tf_video, tf_caption, tf_caption_mask):
    val_data = current_val_data
    val_captions = val_data['caption'].values
    val_captions = map(lambda x: x.replace('.', ''), val_captions)
    val_captions = map(lambda x: x.replace(',', ''), val_captions)
    
    combine_features = load_flickr30k_features if cfg.id == "Flickr30k" else load_msr_vtt_features 
    
    loss_on_validation = []
    for start,end in zip(
                range(0, len(val_data), cfg.batch_size),
                range(cfg.batch_size, len(val_data)+1, cfg.batch_size)): #during every epoch we are discarding incomplete batch in the end
            
        current_batch = val_data[start:end]
        current_videos = current_batch['video_path'].values

        current_feats = np.zeros((cfg.batch_size, cfg.n_frame_step, cfg.dim_image))
        current_feats_vals = map(lambda vid: combine_features(vid), current_videos) 
            
        for ind,feat in enumerate(current_feats_vals):
            current_feats[ind][:len(current_feats_vals[ind])] = feat 
            

        current_captions = current_batch['caption'].values
        current_caption_ind = map(lambda cap: [wordtoix[word] for word in cap.lower().split(' ')[:cfg.n_lstm_step - 1]
                                                   if word in wordtoix],
                                      current_captions)
        
        current_caption_matrix = np.zeros((cfg.batch_size, cfg.n_lstm_step))
        current_caption_masks = np.zeros((cfg.batch_size, cfg.n_lstm_step))
        for ind, row in enumerate(current_caption_masks):
            valid_length = len(current_caption_ind[ind])
            row[:valid_length] = 1
            current_caption_matrix[ind, :valid_length] = current_caption_ind[ind]
        
        loss_val = sess.run(tf_loss,
                feed_dict={
                    tf_video: current_feats,
                    tf_caption: current_caption_matrix,
                    tf_caption_mask: current_caption_masks
                    })
        loss_on_validation.append(loss_val)
    return np.mean(loss_on_validation)
            
        
def train():
    if not os.path.exists(cfg.model_path):
        os.makedirs(cfg.model_path)
    
    print cfg.model_path
    f = open(cfg.model_path + "loss", "a", 1)
    f.write("Checkpoint\tTrain loss\tValidation loss\n")
    
    if cfg.id == "Flickr30k":
        train_data, val_data, _ = get_flickr30k_data(cfg)
    elif cfg.id == "MSR-VTT":
        train_data, val_data, _ = get_msr_vtt_data(cfg)
    
    #FIXME add validation data vocabulary
    captions = train_data['caption'].values
    captions = map(lambda x: x.replace('.', ''), captions)
    captions = map(lambda x: x.replace(',', ''), captions)
    wordtoix, ixtoword, bias_init_vector = preProBuildWordVocab(captions, word_count_threshold=cfg.word_count_threshold)
    
    combine_features = load_flickr30k_features if cfg.id == "Flickr30k" else load_msr_vtt_features

    np.save(cfg.vocab_path + 'ixtoword', ixtoword)
    with open(cfg.vocab_path + 'wordtoix.pkl', 'wb') as outfile:
        pkl.dump(wordtoix, outfile)
    
    sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))
    
    with tf.variable_scope(tf.get_variable_scope()):
        model_train = s2vt(dim_image=cfg.dim_image,
                           n_words=len(ixtoword),
                           dim_hidden=cfg.dim_hidden,
                           batch_size=cfg.batch_size,
                           n_frame_steps=cfg.n_frame_step,
                           n_lstm_steps=cfg.n_lstm_step,
                           dim_word_emb = cfg.dim_word_emb,
                           cell_clip = cfg.cell_clip,
                           forget_bias = cfg.forget_bias,
                           input_keep_prob = cfg.input_keep_prob,
                           output_keep_prob = cfg.output_keep_prob,
                           bias_init_vector=bias_init_vector)
    
        tf_loss, tf_video, tf_caption, tf_caption_mask, _ = model_train.build_model("training")

    with tf.variable_scope(tf.get_variable_scope(), reuse=False):
        train_op = tf.train.AdamOptimizer(cfg.learning_rate).minimize(tf_loss)
    
    saver = tf.train.Saver(max_to_keep=cfg.max_to_keep)
    sess.run(tf.global_variables_initializer())
    
    model_counter = 0
    val_loss = None
    for epoch in range(cfg.n_epochs):
        index = list(train_data.index)
        np.random.shuffle(index)
        train_data = train_data.ix[index]
        
        current_train_data = train_data
        
        total_loss = 0
        saving_schedule = []
        loss_accumulator = []
        
        step_size =  (int(len(current_train_data) * cfg.save_every_n_epoch) // cfg.batch_size ) * cfg.batch_size
        saving_schedule = range(0, len(current_train_data) - step_size, step_size)
        print saving_schedule
        
        for start,end in zip(
                range(0, len(current_train_data), cfg.batch_size),
                range(cfg.batch_size, len(current_train_data)+1, cfg.batch_size)):   
            current_batch = current_train_data[start:end]                           
            current_videos = current_batch['video_path'].values

            current_feats = np.zeros((cfg.batch_size, cfg.n_frame_step, cfg.dim_image))
            current_feats_vals = map(lambda vid: combine_features(vid), current_videos)
            
            
            for ind,feat in enumerate(current_feats_vals):
                current_feats[ind][:len(current_feats_vals[ind])] = feat 
                
            current_captions = current_batch['caption'].values
            current_caption_ind = map(lambda cap: [wordtoix[word] for word in cap.lower().split(' ')[:cfg.n_lstm_step - 1]
                                                   if word in wordtoix],
                                      current_captions)
            
            current_caption_matrix = np.zeros((cfg.batch_size, cfg.n_lstm_step))
            current_caption_masks = np.zeros((cfg.batch_size, cfg.n_lstm_step))
            for ind, row in enumerate(current_caption_masks):
                valid_length = len(current_caption_ind[ind])
                row[:valid_length+1] = 1                                                #forces to predict <EOS> = 0
                current_caption_matrix[ind, :valid_length] = current_caption_ind[ind]
 
            _, train_loss = sess.run(
                    [train_op, tf_loss],
                    feed_dict={
                        tf_video: current_feats,
                        tf_caption: current_caption_matrix,
                        tf_caption_mask: current_caption_masks
                        })
            total_loss += train_loss
            loss_accumulator.append(train_loss)
                
            output_progress(end, len(current_train_data), train_loss)
            
            if start in saving_schedule:
                print start
                train_loss = np.mean(loss_accumulator[-5:])
                val_loss = get_validation_loss(sess, 
                                               val_data.groupby('video_id').apply(lambda x: x.iloc[np.random.choice(len(x))]),
                                               wordtoix, tf_loss, tf_video,
                                               tf_caption, tf_caption_mask)
                f.write(str(model_counter) + "\t" +  str(train_loss) +"\t" + str(val_loss) + "\n")
                sys.stdout.flush()
                saver.save(sess, os.path.join(cfg.model_path, 'model'), global_step=model_counter)
                model_counter+=1
                
        output_progress(end, len(current_train_data), np.mean(loss_accumulator[-5:]))
        print " Done. Validation loss = " + str(val_loss) 
        
        

def convert_data_to_coco_scorer_format(data_frame):
    gts = {}
    non_ascii_count = 0
    for row in zip(data_frame["caption"], data_frame["video_id"]):
        try:
            row[0].encode('ascii', 'ignore').decode('ascii')
        except UnicodeDecodeError:
            non_ascii_count+=1
            continue
        if row[1] in gts:
            gts[row[1]].append({u'image_id': row[1], u'cap_id': len(gts[row[1]]), u'caption':row[0].encode('ascii', 'ignore').decode('ascii')})
        else:
            gts[row[1]] = []
            gts[row[1]].append({u'image_id': row[1], u'cap_id': len(gts[row[1]]), u'caption':row[0].encode('ascii', 'ignore').decode('ascii')})
    if non_ascii_count:
        print "=" * 20 + "\n" + "non-ascii: " + str(non_ascii_count) + "\n" + "=" * 20
    return gts


def test(saved_model=''):
    scorer = COCOScorer()
    ixtoword = pd.Series(np.load(cfg.vocab_path + 'ixtoword.npy').tolist())
    combine_features = load_flickr30k_features if cfg.id == "Flickr30k" else load_msr_vtt_features
    
    model = s2vt(dim_image=cfg.dim_image,
                 n_words=len(ixtoword),
                 dim_hidden=cfg.dim_hidden,
                 batch_size=cfg.batch_size,
                 n_frame_steps=cfg.n_frame_step,
                 n_lstm_steps=cfg.n_lstm_step,
                 dim_word_emb = cfg.dim_word_emb,
                 cell_clip = cfg.cell_clip,
                 forget_bias = cfg.forget_bias,
                 input_keep_prob = cfg.input_keep_prob,
                 output_keep_prob = cfg.output_keep_prob,
                 bias_init_vector=None)
    
    _, video_tf, caption_tf, _, _ = model.build_model("inference")
    session = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))
    saver = tf.train.Saver()
    saver.restore(session, saved_model)
        
    if cfg.id == "Flickr30k":
        _, _, test_data = get_flickr30k_data(cfg)
    elif cfg.id == "MSR-VTT":
        _, _, test_data = get_msr_vtt_data(cfg)
    
    splits = []
    
    splits.append((test_data['video_path'].unique(), test_data))
    results = []
    for split, gt_dataframe in splits:
        gts = convert_data_to_coco_scorer_format(gt_dataframe)
        samples = {}
        for start,end in zip(
                    range(0, len(split), cfg.batch_size),
                    range(cfg.batch_size, len(split) + cfg.batch_size, cfg.batch_size)):
        
            current_batch = split[start:end]
            current_feats = np.zeros((cfg.batch_size, cfg.n_frame_step, cfg.dim_image))
            current_feats_vals = [combine_features(vid) for vid in current_batch] 
            
            for ind,feat in enumerate(current_feats_vals):
                current_feats[ind][:len(current_feats_vals[ind])] = feat
            
            generated_word_index = session.run(caption_tf, feed_dict={video_tf:current_feats})
            generated_word_index = np.asarray(generated_word_index).transpose()
            periods = np.argmax(generated_word_index == 0, axis=1) + 1
            periods[periods == 0] = cfg.n_lstm_step      #take the whole sequence if a period was not produced
            for i in range(len(current_batch)):
                generated_sentence = ' '.join(ixtoword[generated_word_index[i, :periods[i]-1]])
                video_id = current_batch[i].split("/")[-1].split("_")[0] #+ ".jpg"
                samples[video_id] = [{u'image_id': video_id, u'caption': generated_sentence}]
                
        with suppress_stdout_stderr():
            valid_score = scorer.score(gts, samples, samples.keys())
        results.append(valid_score)
        print valid_score
    
    print len(samples)            
    if not os.path.exists(cfg.results_path):
        os.makedirs(cfg.results_path)
    
    with open(cfg.results_path + "scores.txt", 'a') as scores_table:
            scores_table.write(json.dumps(results[0]) + "\n") 
    with open(cfg.results_path + saved_model.split("/")[-1] + ".json", 'w') as prediction_results:
        json.dump({"predictions": samples, "scores": valid_score}, prediction_results)


def main(args):
    global cfg
    if args.dataset == "Flickr30k":
        cfg = flickr_cfg()
    elif args.dataset == "MSR-VTT":
        cfg = msr_vtt_cfg()
    else:
        print "Unknown dataset"
        exit(1)
    
    if args.train_stage:
        train()
    else:
        test(saved_model=cfg.model_path + 'model-' + str(args.checkpoint))
    
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script to train a model for movie description')
    
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--train', dest='train_stage', action='store_true',  
                        help='Training')
    group.add_argument('--test', dest='train_stage', action='store_false', 
                        help='Testing')
    parser.add_argument('--checkpoint', dest='checkpoint', type = int, default = -1,
                        help='Provide a number of the saved model to run testing only on one snapshot')
    parser.add_argument("--dataset", dest='dataset', type=str,
                        help='Specify one from {Flickr30k, MSR-VTT}')
    parser.add_argument("--gpu", dest='gpu', type=str, required=False,
                        help='Set CUDA_VISIBLE_DEVICES environment variable, optional')
    
    args = parser.parse_args()
    if args.gpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    
    if not args.dataset:
        parser.print_help()
        exit(1)

    if not args.train_stage:
        if args.checkpoint is None:
            parser.print_help()
            exit(1)
    
    main(args)