videollama2/inference_audio.py

import os
import json
import math
import argparse
import warnings
import traceback
from tqdm import tqdm

from torch.utils.data import Dataset, DataLoader

import sys
sys.path.append('./')
from videollama2 import model_init, mm_infer
from videollama2.utils import disable_torch_init

# NOTE: Ignore TypedStorage warning, which refers to this link~(https://github.com/pytorch/pytorch/issues/97207#issuecomment-1494781560)
warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')


def split_list(lst, n):
    """Split a list into n (roughly) equal-sized chunks"""
    chunk_size = math.ceil(len(lst) / n)  # integer division
    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]


def get_chunk(lst, n, k):
    chunks = split_list(lst, n)
    return chunks[k]


class ClothoAQADataset(Dataset):

    audoi_formats = ['.wav', '.flac']

    def __init__(self, questions, processor):
        self.questions = questions
        self.processor = processor

    def __len__(self):
        return len(self.questions)
    
    def __getitem__(self, idx):
        sample = self.questions[idx]

        audio_path  = sample['audio']
        question    = sample['conversations'][0]["value"]
        wrapped_question = f"Question: {question}\nAnswer the question using a single word."
        question_id = sample['id']
        answer      = sample['conversations'][1]["value"]

        audio_tensor = self.processor(audio_path)

        return {
            'audio':       audio_tensor,
            'audio_name':  audio_path.split("/")[-1],
            'question':    wrapped_question,
            'question_id': question_id,
            'answer':      answer,
        }

class ClothoDataset(Dataset):

    audoi_formats = ['.wav', '.flac']

    def __init__(self, questions, processor):
        self.questions = questions
        self.processor = processor

    def __len__(self):
        return len(self.questions)
    
    def __getitem__(self, idx):
        sample = self.questions[idx]

        audio_path  = sample['audio']
        wrapped_question = f"Describe the audio."
        question_id = audio_path.split("/")[-1]
        answer      = sample['captions']

        audio_tensor = self.processor(audio_path)

        return {
            'audio':       audio_tensor,
            'audio_name':  audio_path.split("/")[-1],
            'question':    wrapped_question,
            'question_id': question_id,
            'answer':      answer,
        }

class TUT2017Dataset(Dataset):

    audoi_formats = ['.wav', '.flac']

    def __init__(self, questions, processor):
        self.questions = questions
        self.processor = processor

    def __len__(self):
        return len(self.questions)
    
    def __getitem__(self, idx):
        sample = self.questions[idx]

        audio_path  = sample['audio']
        wrapped_question = f"Question: Identify the sound event in the audio.\nOptions:\n(A) beach\n(B) bus\n(C) cafe or restaurant\n(D) car\n(E) city center\n(F) forest path\n(G) grocery store\n(H) home\n(I) library\n(J) metro station\n(K) office\n(L) park\n(M) residential area\n(N) train\n(O) tram\n.Answer with the option's letter from the given choices directly and only give the best option."
        question_id = audio_path.split("/")[-1]
        answer      = sample['gt']

        audio_tensor = self.processor(audio_path)

        return {
            'audio':       audio_tensor,
            'audio_name':  audio_path.split("/")[-1],
            'question':    wrapped_question,
            'question_id': question_id,
            'answer':      answer,
        }

class VocalSoundDataset(Dataset):

    audoi_formats = ['.wav', '.flac']

    def __init__(self, questions, processor):
        self.questions = questions
        self.processor = processor

    def __len__(self):
        return len(self.questions)
    
    def __getitem__(self, idx):
        sample = self.questions[idx]

        audio_path  = sample['audio']
        wrapped_question = f"Identify the human sound in the audio.\nOptions:\n(A) Laughter\n(B) Sigh\n(C) Cough\n(D) Throat clearing\n(E) Sneeze\n(F) Sniff\n.Answer with the option's letter from the given choices directly and only give the best option."
        question_id = audio_path.split("/")[-1]
        answer      = sample['gt']

        audio_tensor = self.processor(audio_path)

        return {
            'audio':       audio_tensor,
            'audio_name':  audio_path.split("/")[-1],
            'question':    wrapped_question,
            'question_id': question_id,
            'answer':      answer,
        }

class AIRDataset(Dataset):

    audoi_formats = ['.wav', '.flac']

    def __init__(self, questions, processor):
        self.questions = questions
        self.processor = processor

    def __len__(self):
        return len(self.questions)
    
    def __getitem__(self, idx):
        sample = self.questions[idx]

        audio_path = sample['audio']
        wrapped_question = sample['query']
        question_id = sample['id']
        answer = sample['answer']

        audio_tensor = self.processor(audio_path)

        return {
            'audio':       audio_tensor,
            'audio_name':  audio_path.split("/")[-1],
            'question':    wrapped_question,
            'question_id': question_id,
            'answer':      answer,
        }


def collate_fn(batch):
    vid  = [x['audio'] for x in batch]
    v_id = [x['audio_name'] for x in batch]
    qus  = [x['question'] for x in batch]
    qid  = [x['question_id'] for x in batch]
    ans  = [x['answer'] for x in batch]
    return vid, v_id, qus, qid, ans


def run_inference(args):
    disable_torch_init()

    # Initialize the model
    model, processor, tokenizer = model_init(args.model_path)
    model.model.vision_tower = None

    assert args.batch_size == 1, "Batch size must be 1 for inference"
    if args.dataset == "clothoAQA":
        gt_questions = json.load(open(args.question_file, "r"))
        gt_questions = get_chunk(gt_questions, args.num_chunks, args.chunk_idx)
        dataset = ClothoAQADataset(gt_questions, processor['audio'])
    elif args.dataset == "clotho":
        import csv
        gt_questions = []
        with open(args.question_file, mode='r', encoding='utf-8') as file:
            reader = csv.reader(file)
            header = next(reader) # remove header
            for row in reader:
                gt_questions.append({
                    "audio": os.path.join(args.video_folder, row[0]),
                    "captions": row[1:]
                })
        gt_questions = get_chunk(gt_questions, args.num_chunks, args.chunk_idx)
        dataset = ClothoDataset(gt_questions, processor['audio'])
    elif args.dataset == "TUT2017":
        gt_questions = []
        with open(args.question_file, "r") as fp:
            for x in fp.readlines():
                gt_questions.append(json.loads(x))
                gt_questions[-1]["audio"] = os.path.join(args.video_folder, gt_questions[-1]["audio"])
        gt_questions = get_chunk(gt_questions, args.num_chunks, args.chunk_idx)
        dataset = TUT2017Dataset(gt_questions, processor['audio'])
    elif args.dataset == "vocalsound":
        gt_questions = []
        with open(args.question_file, "r") as fp:
            for x in fp.readlines():
                gt_questions.append(json.loads(x))
                gt_questions[-1]["audio"] = os.path.join(args.video_folder, gt_questions[-1]["audio"].split("/")[-1])
        gt_questions = get_chunk(gt_questions, args.num_chunks, args.chunk_idx)
        dataset = VocalSoundDataset(gt_questions, processor['audio'])
    elif args.dataset == "AIR":
        gt_answer = {x["uniq_id"]: x for x in json.load(open(args.answer_file, "r"))}
        
        gt_questions = []
        with open(args.question_file, "r") as fp:
            for x in fp.readlines():
                gt_questions.append(json.loads(x))
                gt_questions[-1]["answer"] = gt_answer[gt_questions[-1]["id"]]["answer_gt"]
        gt_questions = get_chunk(gt_questions, args.num_chunks, args.chunk_idx)
        dataset = AIRDataset(gt_questions, processor['audio'])
    else:
        raise NotImplementedError

    dataloader = DataLoader(dataset, shuffle=False, batch_size=args.batch_size, num_workers=args.num_workers, collate_fn=collate_fn)

    answer_file = os.path.join(args.output_file)
    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    ans_file = open(answer_file, "w")

    # Iterate over each sample in the ground truth file
    for i, (audio_tensors, audio_names, questions, question_ids, answers) in enumerate(tqdm(dataloader)):
        audio_tensor = audio_tensors[0]
        audio_name   = audio_names[0]
        question     = questions[0]
        question_id  = question_ids[0]
        answer       = answers[0]

        # question = question + '\n' + 'Answer the question using a single word or a short phrase with multiple words.'

        try:
            output = mm_infer(
                audio_tensor,
                question,
                model=model,
                tokenizer=tokenizer,
                modal='audio',
                do_sample=False,
            )
        except:
            traceback.print_exc()
            output = "error"

        sample_set = {'id': question_id, 'question': question, 'answer': answer, 'pred': output}
        ans_file.write(json.dumps(sample_set) + "\n")

    ans_file.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('--model-path', help='', required=True)
    parser.add_argument('--video-folder', help='Directory containing video files.', required=True)
    parser.add_argument('--question-file', help='Path to the ground truth file containing question.', required=True)
    parser.add_argument('--answer-file', help='Path to the ground truth file containing answers.', required=False)
    parser.add_argument('--output-file', help='Directory to save the model results JSON.', required=True)
    parser.add_argument("--num-chunks", type=int, default=1)
    parser.add_argument("--chunk-idx", type=int, default=0)
    parser.add_argument("--device", type=str, required=False, default='cuda:0')
    parser.add_argument("--batch-size", type=int, required=False, default=1)
    parser.add_argument("--num-workers", type=int, required=False, default=8)
    parser.add_argument("--dataset", type=str, required=True)
    args = parser.parse_args()

    run_inference(args)