IMAX/eval_reg.py at main · MSIIP/IMAX · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import os
import json
import numpy as np
from tqdm import tqdm
from torch.utils.data import DataLoader

from imax.common.config import Config
from imax.common.eval_utils import prepare_texts, init_model, eval_parser
from imax.conversation.conversation import Conversation

from evaluate_metrics import calculate_bleu, calculate_equal_acc
from imax.datasets.datasets.mimic_rec_reg_dataset import Mimic_Reg_EvalData

def evaluate_reg(model_predict):

    bleu_scores = calculate_bleu(model_predict)
    acc_scores = calculate_equal_acc(model_predict)

    return {
        "bleu_scores": bleu_scores,
        "acc_score": acc_scores,
    }

def list_of_str(arg):
    return list(map(str, arg.split(',')))

parser = eval_parser()
parser.add_argument("--dataset", type=list_of_str, help="dataset to evaluate")
args = parser.parse_args()

cfg = Config(args)


model, vis_processor, text_processor = init_model(args)
bos_token = model.llm_tokenizer.bos_token
conv_temp = Conversation(
    system="",
    roles=(f"{bos_token}[INST] ", " [/INST]"),
    messages=[],
    sep="",
    offset=2,
).copy()
model.eval()
cfg_save_path = cfg.run_cfg.save_path
save_path = os.path.join(cfg_save_path, str(args.dataset[0]))

if 'mimic_reg' in args.dataset:
    eval_file_path = cfg.evaluation_datasets_cfg["mimic_reg"]["eval_file_path"]
    img_path = cfg.evaluation_datasets_cfg["mimic_reg"]["img_path"]
    batch_size = cfg.evaluation_datasets_cfg["mimic_reg"]["batch_size"]
    max_new_tokens = cfg.evaluation_datasets_cfg["mimic_reg"]["max_new_tokens"]

    data = Mimic_Reg_EvalData(vis_processor, text_processor, img_path, [eval_file_path])
    eval_dataloader = DataLoader(data, batch_size=batch_size, shuffle=False)
    model_predict = []

    for image_paths, images, questions, answers_gt in tqdm(eval_dataloader):
        texts = prepare_texts(questions, conv_temp)
        answers_pred = model.generate(images, texts, questions=questions, max_new_tokens=max_new_tokens, do_sample=False)


        for image_path, question, answer_gt , answer_pred \
            in zip(image_paths, questions, answers_gt, answers_pred):
            result = dict()
            answer_pred = answer_pred.lower().replace('<unk>','').strip()

            result['image_path'] = image_path
            result['question'] = question
            result['answer_gt'] = answer_gt
            result['answer_pred'] = answer_pred

            model_predict.append(result)

    # save the result
    if not os.path.exists(save_path):
        os.makedirs(save_path, exist_ok=True)
    file_save_path= os.path.join(save_path, "mimic_reg.json")
    with open(file_save_path,'w') as f:
        json.dump(model_predict, f)
    print("save the result to {}".format(file_save_path))

    # evaluate the result
    metrics = evaluate_reg(model_predict)
    metric_save_path = os.path.join(save_path, "metric.json")
    with open(metric_save_path, 'w') as f:
        json.dump(metrics, f, sort_keys=True)
    print("save the metrics to {}".format(metric_save_path))
    print("metrics: {}".format(metrics))