gtad_inference.py

import os
import math
import numpy as np
import pandas as pd
import torch.nn.parallel

from gtad_lib import opts
from gtad_lib.models import GTAD
from gtad_lib.dataset import VideoDataSet

if __name__ == "__main__":
    opt = opts.parse_opt()
    opt = vars(opt)
    if not os.path.exists(opt["output"] + "/results"):
        os.makedirs(opt["output"] + "/results")

    model = GTAD(opt)
    model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
    checkpoint = torch.load(opt["output"] + "/GTAD_best.pth.tar")
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()

    test_loader = torch.utils.data.DataLoader(
        VideoDataSet(opt, subset="validation", mode="inference"),
        batch_size=1,
        shuffle=False,
        num_workers=8,
        pin_memory=True,
        drop_last=False,
    )

    print("Inference start")
    with torch.no_grad():
        for idx, input_data in test_loader:
            video_name = test_loader.dataset.video_list[idx[0]]
            offset = min(test_loader.dataset.data["indices"][idx[0]])
            video_name = video_name + "_{}".format(math.floor(offset / 250))
            input_data = input_data.cuda()

            # forward pass
            confidence_map, _, _ = model(input_data)

            clr_confidence = (confidence_map[0][1]).detach().cpu().numpy()
            reg_confidence = (confidence_map[0][0]).detach().cpu().numpy()

            # enumerate sub-graphs as proposals
            new_props = []
            for idx in range(opt["max_duration"]):
                for jdx in range(opt["temporal_scale"]):
                    start_index = jdx
                    end_index = start_index + idx + 1
                    if end_index < opt["temporal_scale"]:
                        xmin = (
                            start_index * opt["skip_videoframes"] + offset
                        )  # map [0,99] to frames
                        xmax = end_index * opt["skip_videoframes"] + offset
                        clr_score = clr_confidence[idx, jdx]  # 64, 128
                        reg_score = reg_confidence[idx, jdx]
                        new_props.append([xmin, xmax, clr_score, reg_score])
            new_props = np.stack(new_props)

            col_name = ["xmin", "xmax", "clr_score", "reg_socre"]
            new_df = pd.DataFrame(new_props, columns=col_name)
            new_df.to_csv(
                opt["output"] + "/results/" + video_name + ".csv", index=False
            )

    print("Inference finished")