app.py

import numpy as np
# import gradio
import torch
from transformers import BertTokenizer
import argparse
import gradio as gr
import time

from modules.tokenization_clip import SimpleTokenizer as ClipTokenizer
from modules.modeling import BirdModel

show_num = 10
max_words = 32
video_path_zh = "features/Chinese_batch_visual_output_list.npy"
frame_path_zh = "features/Chinese_batch_frame_output_list.npy"
video_fea_zh = np.load(video_path_zh)
video_fea_zh = torch.from_numpy(video_fea_zh)
frame_fea_zh = np.load(frame_path_zh)
frame_fea_zh = torch.from_numpy(frame_fea_zh)

video_path_en = "features/English_batch_visual_output_list.npy"
frame_path_en = "features/English_batch_frame_output_list.npy"
video_fea_en = np.load(video_path_en)
video_fea_en = torch.from_numpy(video_fea_en)
frame_fea_en = np.load(frame_path_en)
frame_fea_en = torch.from_numpy(frame_fea_en)

test_path = "test_list.txt"
# video_dir = "test1500_400_400/"
video_dir = "test1500/"

with open(test_path, 'r', encoding='utf8') as f_list:
    lines = f_list.readlines()
    video_ids = [itm.strip() + ".mp4" for itm in lines]


def get_videoname(idx):
    videoname = []
    videopath = []
    for i in idx:
        videoname.append(video_ids[i])
        path = video_dir + video_ids[i]
        videopath.append(path)
    return videoname, videopath


def get_text(caption, tokenizer):
    # tokenize word
    words = tokenizer.tokenize(caption)

    # add cls token
    words = ["<|startoftext|>"] + words
    total_length_with_CLS = max_words - 1
    if len(words) > total_length_with_CLS:
        words = words[:total_length_with_CLS]

    # add end token
    words = words + ["<|endoftext|>"]

    # convert token to id according to the vocab
    input_ids = tokenizer.convert_tokens_to_ids(words)

    # add zeros for feature of the same length
    input_mask = [1] * len(input_ids)
    while len(input_ids) < max_words:
        input_ids.append(0)
        input_mask.append(0)

    # ensure the length of feature to be equal with max words
    assert len(input_ids) == max_words
    assert len(input_mask) == max_words
    pairs_text = np.array(input_ids).reshape(-1, max_words)
    pairs_text = torch.from_numpy(pairs_text)
    pairs_mask = np.array(input_mask).reshape(-1, max_words)
    pairs_mask = torch.from_numpy(pairs_mask)

    return pairs_text, pairs_mask


def get_args(description='Retrieval Task'):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--do_pretrain", action='store_true', help="Whether to run training.")
    parser.add_argument("--do_train", action='store_true', help="Whether to run training.")
    parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.")
    parser.add_argument("--do_params", action='store_true', help="text the params of the model.")
    parser.add_argument("--use_frame_fea", action='store_true', help="whether use frame feature matching text")
    parser.add_argument('--task', type=str, default="retrieval", choices=["retrieval_VT", "retrieval"],
                        help="choose downstream task.")
    parser.add_argument('--dataset', type=str, default="bird", choices=["bird", "msrvtt", "vatex", "msvd"],
                        help="choose dataset.")
    parser.add_argument('--num_thread_reader', type=int, default=1, help='')
    parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate')
    parser.add_argument('--text_lr', type=float, default=0.00001, help='text encoder learning rate')
    parser.add_argument('--epochs', type=int, default=20, help='upper epoch limit')
    parser.add_argument('--batch_size', type=int, default=256, help='batch size')
    parser.add_argument('--batch_size_val', type=int, default=3500, help='batch size eval')
    parser.add_argument('--lr_decay', type=float, default=0.9, help='Learning rate exp epoch decay')
    parser.add_argument('--weight_decay', type=float, default=0.2, help='Learning rate exp epoch decay')
    parser.add_argument('--n_display', type=int, default=100, help='Information display frequence')
    parser.add_argument('--seed', type=int, default=42, help='random seed')
    parser.add_argument('--max_words', type=int, default=32, help='')
    parser.add_argument('--max_frames', type=int, default=12, help='')
    parser.add_argument('--top_frames', type=int, default=3, help='')
    parser.add_argument('--frame_sample', type=str, default="uniform", choices=["uniform", "random", "uniform_random"],
                        help='frame sample strategy')
    parser.add_argument('--frame_sample_len', type=str, default="fix", choices=["dynamic", "fix"],
                        help='use dynamic frame length of fix frame length')
    parser.add_argument('--language', type=str, default="chinese", choices=["chinese", "english"],
                        help='language for text encoder')
    parser.add_argument('--use_temp', action='store_true', help='whether to use temporal transformer')

    parser.add_argument("--logdir", default=None, type=str, required=False, help="log dir for tensorboardX writer")
    parser.add_argument("--cross_model", default="cross-base", type=str, required=False, help="Cross module")
    parser.add_argument("--pretrained_text", default="hfl/chinese-roberta-wwm-ext", type=str, required=False, help="pretrained_text")
    parser.add_argument("--init_model", default=None, type=str, required=False, help="Initial model.")
    parser.add_argument("--warmup_proportion", default=0.1, type=float,
                        help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument('--n_gpu', type=int, default=1, help="Changed in the execute process.")

    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument('--enable_amp', action='store_true', help="whether to use pytorch amp")

    parser.add_argument("--world_size", default=0, type=int, help="distribted training")
    parser.add_argument("--local_rank", default=0, type=int, help="distribted training")
    parser.add_argument("--rank", default=0, type=int, help="distribted training")
    parser.add_argument('--coef_lr', type=float, default=1., help='coefficient for bert branch.')

    args = parser.parse_args()

    # Check paramenters
    args.do_eval = True
    args.use_frame_fea = True
    args.use_temp = True

    return args


def init_model(language):
    time1 = time.time()
    args = get_args()
    args.language = language
    if language == "chinese":
        model_path = "models/Chinese_vatex.bin"
        tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext")
    elif language == "english":
        model_path = "models/English_vatex.bin"
        tokenizer = ClipTokenizer()
    else:
        raise Exception("language should be Chinese or English!")
    model_state_dict = torch.load(model_path, map_location='cpu')
    cross_model = "cross-base"
    model = BirdModel.from_pretrained(cross_model, state_dict=model_state_dict, task_config=args)
    device = torch.device("cpu")
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    print("language={}".format(language))
    print("init model time: {}".format(time.time() - time1))
    print("device:{}".format(device))
    return model, tokenizer


model_zh, tokenizer_zh = init_model(language="chinese")
model_en, tokenizer_en = init_model(language="english")


def t2v_search_zh(text):
    with torch.no_grad():
        time1 = time.time()
        text_ids, text_mask = get_text(text, tokenizer_zh)
        print("get_text time: {}".format(time.time() - time1))
        time1 = time.time()
        text_fea_zh = model_zh.text_encoder(text_ids, text_mask)
        print("text_encoder time: {}".format(time.time() - time1))
        # print("text_fea.shape:{}".format(text_fea.shape))
        # print("video_fea.shape:{}".format(video_fea.shape))
        # print("frame_fea.shape:{}".format(frame_fea.shape))
        time1 = time.time()
        sim_video = model_zh.loose_similarity(text_fea_zh, video_fea_zh)
        # print("sim_video.shape:{}".format(sim_video.shape))
        sim_frame = model_zh.loose_similarity(text_fea_zh, frame_fea_zh)
        # print("sim_frame.shape:{}".format(sim_frame.shape))
        sim_frame = torch.topk(sim_frame, k=model_zh.top_frames, dim=1)[0]
        sim_frame = torch.mean(sim_frame, dim=1)
        sim = sim_video + sim_frame
        value, index = sim.topk(show_num, dim=0, largest=True, sorted=True)
        # value, index = sim_video.topk(show_num, dim=0, largest=True, sorted=True)
        print("calculate_similarity time: {}".format(time.time() - time1))
        print("value:{}".format(value))
        print("index:{}".format(index))
        videoname, videopath = get_videoname(index)
        print("videoname:{}".format(videoname))
        print("videopath:{}".format(videopath))
        return videopath


def t2v_search_en(text):
    with torch.no_grad():
        time1 = time.time()
        text_ids, text_mask = get_text(text, tokenizer_en)
        print("get_text time: {}".format(time.time() - time1))
        time1 = time.time()
        text_fea_en = model_en.text_encoder(text_ids, text_mask)
        print("text_encoder time: {}".format(time.time() - time1))
        # print("text_fea.shape:{}".format(text_fea.shape))
        # print("video_fea.shape:{}".format(video_fea.shape))
        # print("frame_fea.shape:{}".format(frame_fea.shape))
        time1 = time.time()
        sim_video = model_en.loose_similarity(text_fea_en, video_fea_en)
        # print("sim_video.shape:{}".format(sim_video.shape))
        sim_frame = model_en.loose_similarity(text_fea_en, frame_fea_en)
        # print("sim_frame.shape:{}".format(sim_frame.shape))
        sim_frame = torch.topk(sim_frame, k=model_en.top_frames, dim=1)[0]
        sim_frame = torch.mean(sim_frame, dim=1)
        sim = sim_video + sim_frame
        value, index = sim.topk(show_num, dim=0, largest=True, sorted=True)
        # value, index = sim_video.topk(show_num, dim=0, largest=True, sorted=True)
        print("calculate_similarity time: {}".format(time.time() - time1))
        print("value:{}".format(value))
        print("index:{}".format(index))
        videoname, videopath = get_videoname(index)
        print("videoname:{}".format(videoname))
        print("videopath:{}".format(videopath))
        return videopath


def hello_world(name):
    return "hello world, my name is " + name + "!"


def search_demo():
    with gr.Blocks() as demo:
        gr.Markdown("# <center>HMMC中英文本-视频检索</center>")
        demo.title = "HMMC中英文本-视频检索"
        with gr.Tab("中文"):
            with gr.Column(variant="panel"):
                with gr.Row(variant="compact"):
                    input_text = gr.Textbox(
                        label="输入文本",
                        show_label=False,
                        max_lines=1,
                        placeholder="请输入检索文本...",
                    ).style(
                        container=False,
                    )
                    btn = gr.Button("搜索").style(full_width=False)

                with gr.Column(variant="panel", scale=2):
                    with gr.Row(variant="compact"):
                        videos_top = [gr.Video(
                            format="mp4", label="视频 "+str(i+1),
                        ).style(height=300, width=300) for i in range(3)]
                with gr.Column(variant="panel", scale=1):
                    with gr.Row(variant="compact"):
                        videos_rest = [gr.Video(
                            format="mp4", label="视频 "+str(i),
                        ).style(height=150, width=150) for i in range(3, show_num)]

            searched_videos = videos_top + videos_rest
            btn.click(t2v_search_zh, inputs=input_text, outputs=searched_videos)

        with gr.Tab("English"):
            with gr.Column(variant="panel"):
                with gr.Row(variant="compact"):
                    input_text = gr.Textbox(
                        label="input text",
                        show_label=False,
                        max_lines=1,
                        placeholder="Please input text to search...",
                    ).style(
                        container=False,
                    )
                    btn = gr.Button("Search").style(full_width=False)

                with gr.Column(variant="panel", scale=2):
                    with gr.Row(variant="compact"):
                        videos_top = [gr.Video(
                            format="mp4", label="video " + str(i+1),
                        ).style(height=300, width=300) for i in range(3)]
                with gr.Column(variant="panel", scale=1):
                    with gr.Row(variant="compact"):
                        videos_rest = [gr.Video(
                            format="mp4", label="video " + str(i),
                        ).style(height=150, width=150) for i in range(3, show_num)]

            searched_videos = videos_top + videos_rest
            btn.click(t2v_search_en, inputs=input_text, outputs=searched_videos)

    demo.launch(share=True)


if __name__ == '__main__':
    search_demo()
    # text = "两个男人正在随着音乐跳舞，他们正在努力做着macarena舞蹈的动作。"
    # t2v_search(text)