Spaces:
Running on Zero
Running on Zero
File size: 2,400 Bytes
d43892c | 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 | from argparse import ArgumentParser
import utils
import torch
from models.basic_model import CDEvaluator
import os
"""
quick start
sample files in ./samples
save prediction files in the ./samples/predict
"""
def get_args():
# ------------
# args
# ------------
parser = ArgumentParser()
parser.add_argument('--project_name', default='BIT_LEVIR', type=str)
parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
parser.add_argument('--checkpoint_root', default='checkpoints', type=str)
parser.add_argument('--output_folder', default='samples/predict', type=str)
# data
parser.add_argument('--num_workers', default=0, type=int)
parser.add_argument('--dataset', default='CDDataset', type=str)
parser.add_argument('--data_name', default='quick_start', type=str)
parser.add_argument('--batch_size', default=1, type=int)
parser.add_argument('--split', default="demo", type=str)
parser.add_argument('--img_size', default=256, type=int)
# model
parser.add_argument('--n_class', default=2, type=int)
parser.add_argument('--net_G', default='base_transformer_pos_s4_dd8_dedim8', type=str,
help='base_resnet18 | base_transformer_pos_s4_dd8 | base_transformer_pos_s4_dd8_dedim8|')
parser.add_argument('--checkpoint_name', default='best_ckpt.pt', type=str)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
utils.get_device(args)
device = torch.device("cuda:%s" % args.gpu_ids[0]
if torch.cuda.is_available() and len(args.gpu_ids)>0
else "cpu")
args.checkpoint_dir = os.path.join(args.checkpoint_root, args.project_name)
os.makedirs(args.output_folder, exist_ok=True)
log_path = os.path.join(args.output_folder, 'log_vis.txt')
data_loader = utils.get_loader(args.data_name, img_size=args.img_size,
batch_size=args.batch_size,
split=args.split, is_train=False)
model = CDEvaluator(args)
model.load_checkpoint(args.checkpoint_name)
model.eval()
for i, batch in enumerate(data_loader):
name = batch['name']
print('process: %s' % name)
score_map = model._forward_pass(batch)
model._save_predictions()
|