-
Notifications
You must be signed in to change notification settings - Fork 7
/
infer.py
175 lines (161 loc) · 11.4 KB
/
infer.py
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import gzip
import os
import pickle
import random
import sys
import time
import traceback
import nltk
import numpy as np
import stanza
import torch.utils.data
from torch.utils.data import DataLoader
from clinicgen.data.image2text import PretrainedEmbeddings
from clinicgen.data.utils import Data
from clinicgen.eval import GenEval
from clinicgen.log import EpochLog, FileLogger
from clinicgen.models.utils import Models
def main(args):
# Download
if args.nltk_download:
nltk.download('punkt')
if args.stanza_download:
stanza.download('en', processors='tokenize,lemma,pos,ner')
stanza.download('en', package='radiology')
# Set random seeds
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# Embeddings
embeddings, word_idxs = PretrainedEmbeddings.load_embeddings(args.embeddings)
# Data
t = time.time()
hierarchical = Models.hierarchical(args.model)
max_sent = args.max_sent if hierarchical else None
datasets = Data.get_datasets(args.data, args.corpus, word_idxs, args.sentsplitter, args.tokenizer, args.textfilter,
args.tokenfilter, max_sent, args.max_word, multi_image=args.multi_image,
img_mode=args.img_trans, img_augment=args.img_augment, single_test=args.single_test,
cache_data=args.cache_data, section=args.section, anatomy=args.anatomy,
meta=args.splits, exclude_ids=args.exclude_ids, a_labels=args.a_labels,
test_only=True)
nw = 0 if args.cache_data else args.num_workers
val_loader = DataLoader(datasets['validation'], batch_size=args.batch_size, shuffle=False, num_workers=nw,
pin_memory=args.pin_memory)
test_loader = DataLoader(datasets['test'], batch_size=args.batch_size, shuffle=False, num_workers=nw,
pin_memory=args.pin_memory)
print('Data: validation={0}, test={1} (load time {2:.2f}s)'.format(len(val_loader.dataset),
len(test_loader.dataset), time.time() - t))
# reports = []
# for cur_elem in test_loader.dataset:
# reports.append(cur_elem[2])
# import pdb
# pdb.set_trace()
# import pandas as pd
# df = pd.DataFrame({'Report Impression': reports})
# df.to_csv('ifcc_test.csv', index=False)
if args.cider_df is not None:
with gzip.open(args.cider_df) as f:
cider_df = pickle.load(f)
print('CIDEr df: {0}'.format(len(cider_df)))
else:
cider_df = None
# Model configurations
device = 'gpu' if args.cuda else 'cpu'
model = Models.get_model(args.model, embeddings, args.hidden_size, args.max_word, args.max_sent,
multi_image=args.multi_image, multi_merge=args.multi_merge, image_model=args.img_model,
image_pretrained=args.img_pretrained, view_position=args.view_position,
parallel_sent=args.parallel_sent, word_idxs=word_idxs, device=device,
cnnrnnrnn_topic_state=args.cnnrnnrnn_topic_state,
cnnrnnrnn_simple_proj=args.cnnrnnrnn_simple_proj, sat_lstm_dim=args.sat_lstm_dim,
trans_image_pe=args.img_pe, trans_layers=args.trans_layers,
trans_enc_layers=args.trans_enc_layers, trans_layer_norm=args.trans_layer_norm,
m2_memory=args.m2_memory, tienet_labels=args.tienet_labels, verbose=args.verbose)
if device == 'gpu':
model = model.cuda(0)
# Train and test processes
evaluator = GenEval(model, word_idxs, beam_size=args.beam_size, cider_df=cider_df, spice=args.spice,
bert_score=args.bert_score, bert_score_penalty=args.bert_score_penalty, nli=args.nli,
nli_compare=args.nli_comp, nli_label=args.nli_label, nli_neutral_score=args.nli_neutral_score,
nli_prf=args.nli_prf, nli_batch=args.nli_batch, entity_match=args.entity_match,
entity_mode=args.entity_mode, nthreads=args.nthreads, pin_memory=args.pin_memory,
verbose=args.verbose)
with FileLogger(args.out, args.out, model, evaluator, None, None, None, None, mode='all', device=device) as logger:
try:
pbar_vals = {'losses': None, 'val_scores': None, 'test_scores': None}
bep = logger.load_baseline(args.baseline_model)
print('Starting from baseline {0} ({1})'.format(os.path.basename(args.baseline_model), bep + 1))
# Run evaluation on the baseline model
EpochLog.log_datasets(logger, pbar_vals, 0, 0, None, val_loader, test_loader, save=False, progress=True)
logger.save_parameters(args)
except BaseException:
print('Unexpected exception: {0}'.format(traceback.format_exc()))
def parse_args():
parser = argparse.ArgumentParser(description='Infer from a model for report generation')
parser.add_argument('data', type=str, help='A path to clinical data')
parser.add_argument('baseline_model', type=str, help='A baseline model to start from')
parser.add_argument('embeddings', type=str, help='A path to word embeddings')
parser.add_argument('out', type=str, help='An output path')
parser.add_argument('--a-labels', type=str, default=None, help='AReport validity labels')
parser.add_argument('--anatomy', type=str, default=None, help='A specific anatomy to target')
parser.add_argument('--batch-size', type=int, default=24, help='Batch size')
parser.add_argument('--beam-size', type=int, default=4, help='Beam size')
parser.add_argument('--bert-score', type=str, default=None, help='BERTScore model type')
parser.add_argument('--bert-score-penalty', default=False, action='store_true', help='Add a Gaussian penalty to BERTScore')
parser.add_argument('--cache-data', type=str, default=None, help='Cache images and texts to memory and disk')
parser.add_argument('--cider-df', type=str, default=None, help='A path to CIDEr DF')
parser.add_argument('--cnnrnnrnn-mlp-proj', dest='cnnrnnrnn_simple_proj', default=True, action='store_false', help='An MLP visual feature projection for CNNRNNRNN')
parser.add_argument('--cnnrnnrnn-topic-state', default=False, action='store_true', help='Use topic as an initial word LSTM state')
parser.add_argument('--corpus', type=str, default='a', choices=['a', 'flickr30k', 'mimic-cxr', 'open-i'], help='Corpus name')
parser.add_argument('--cuda', default=False, action='store_true', help='Use GPU')
parser.add_argument('--entity-match', type=str, default=None, help='A path to reference entities')
parser.add_argument('--entity-mode',type=str, default='nli-f', help='Entity match mode')
parser.add_argument('--exclude-ids', type=str, default=None, help='IDs to exclude from the data')
parser.add_argument('--hidden-size', type=int, default=512, help='Hidden unit size')
parser.add_argument('--img-no-augment', dest='img_augment', default=True, action='store_false', help='Do not augment images')
parser.add_argument('--img-pe', default=False, action='store_true', help='Add positional encodings for images')
parser.add_argument('--img-trans', type=str, default='pad', choices=['center', 'pad'], help='Image transformation mode')
parser.add_argument('--img-model', type=str, default=None, help='An image model')
parser.add_argument('--img-pretrained', type=str, default=None, help='Pre-trained image model')
parser.add_argument('--iter-sent', dest='parallel_sent', default=True, action='store_false', help='Iteratively process sentences')
parser.add_argument('--m2-memory', type=int, default=40, help='M2 Transformer memory size')
parser.add_argument('--max-sent', type=int, default=1, help='Max sentence num')
parser.add_argument('--max-word', type=int, default=128, help='Max word num')
parser.add_argument('--model', type=str, default='m2trans', choices=['cnnrnnrnn', 'kwl', 'm2trans', 'sat', 'tienet', 'trans', 'trans-s'])
parser.add_argument('--multi-image', type=int, default=1, help='Multi image number')
parser.add_argument('--multi-merge', type=str, default='max', choices=['att', 'max'], help='A merge method for multi images')
parser.add_argument('--nli', type=str, default=None, choices=['mednli', 'mednli-rad'], help='NLI model type')
parser.add_argument('--nli-batch', type=int, default=24, help='NLI batch size')
parser.add_argument('--nli-comp', type=str, default='bert-score', help='NLI comparison method')
parser.add_argument('--nli-label', type=str, default='entailment', choices=['contradiction', 'entailment'], help='NLI score label')
parser.add_argument('--nli-neutral-score', type=float, default=(1.0 / 3), help='An NLI entailment neutral score')
parser.add_argument('--nli-prf', type=str, default='f', choices=['f', 'fh', 'fp', 'p', 'r'], help='NLI metric')
parser.add_argument('--nltk-download', default=False, action='store_true', help='Download NLTK punkt data')
parser.add_argument('--nthreads', type=int, default=2, help='Number of threads')
parser.add_argument('--num-workers', type=int, default=1, help='Number of background workers for data loader')
parser.add_argument('--pin-memory', dest='pin_memory', default=False, action='store_true', help='Use pin-memory on data loaders')
parser.add_argument('--sat-lstm-dim', type=int, default=1000, help='An LSTM dimension for SAT')
parser.add_argument('--section', type=str, default='findings', help='Report section')
parser.add_argument('--seed', type=int, default=1, help='Random seed')
parser.add_argument('--sentsplitter', type=str, default='none', choices=['linebreak', 'nltk', 'none', 'stanford'], help='Sentence splitter name')
parser.add_argument('--single-test', default=False, action='store_true', help='Test with a single image setting')
parser.add_argument('--spice', default=False, action='store_true', help='SPICE evaluation')
parser.add_argument('--splits', type=str, default=None, help='A path to a file defining splits')
parser.add_argument('--stanza-download', default=False, action='store_true', help='Download Stanza clinical model')
parser.add_argument('--textfilter', type=str, default='lower', help='Text filter')
parser.add_argument('--tienet-labels', type=str, default=None, help='TieNet labels')
parser.add_argument('--tokenfilter', type=str, default='none', help='Token filter')
parser.add_argument('--tokenizer', type=str, default='nltk', choices=['nltk', 'none', 'stanford', 'whitespace'], help='Tokenizer name')
parser.add_argument('--trans-enc-layers', type=int, default=None, help='Number of transformer encoder layers')
parser.add_argument('--trans-layers', type=int, default=1, help='Number of transformer layers')
parser.add_argument('--trans-layer-no-norm', dest='trans_layer_norm', default=True, action='store_false', help='Do not Layer normalize visual features in transformer models')
parser.add_argument('--verbose', default=False, action='store_true', help='Verbose outputs')
parser.add_argument('--view-position', default=False, action='store_true', help='Include view position embeddings')
return parser.parse_args()
if __name__ == '__main__':
cdir = os.path.abspath(os.path.dirname(__file__))
sys.path.append(cdir)
args = parse_args()
main(args)