File size: 3,785 Bytes
c614b0f |
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 |
import os
import argparse
from glob import glob
from tqdm import tqdm
import cv2
import torch
from dataset import MyData
from models.birefnet import BiRefNet, BiRefNetC2F
from utils import save_tensor_img, check_state_dict
from config import Config
config = Config()
def inference(model, data_loader_test, pred_root, method, testset, device=0):
model_training = model.training
if model_training:
model.eval()
for batch in tqdm(data_loader_test, total=len(data_loader_test)) if 1 or config.verbose_eval else data_loader_test:
inputs = batch[0].to(device)
# gts = batch[1].to(device)
label_paths = batch[-1]
with torch.no_grad():
scaled_preds = model(inputs)[-1].sigmoid()
os.makedirs(os.path.join(pred_root, method, testset), exist_ok=True)
for idx_sample in range(scaled_preds.shape[0]):
res = torch.nn.functional.interpolate(
scaled_preds[idx_sample].unsqueeze(0),
size=cv2.imread(label_paths[idx_sample], cv2.IMREAD_GRAYSCALE).shape[:2],
mode='bilinear',
align_corners=True
)
save_tensor_img(res, os.path.join(os.path.join(pred_root, method, testset), label_paths[idx_sample].replace('\\', '/').split('/')[-1])) # test set dir + file name
if model_training:
model.train()
return None
def main(args):
# Init model
device = config.device
if args.ckpt_folder:
print('Testing with models in {}'.format(args.ckpt_folder))
else:
print('Testing with model {}'.format(args.ckpt))
if config.model == 'BiRefNet':
model = BiRefNet(bb_pretrained=False)
elif config.model == 'BiRefNetC2F':
model = BiRefNetC2F(bb_pretrained=False)
weights_lst = sorted(
glob(os.path.join(args.ckpt_folder, '*.pth')) if args.ckpt_folder else [args.ckpt],
key=lambda x: int(x.split('epoch_')[-1].split('.pth')[0]),
reverse=True
)
for testset in args.testsets.split('+'):
print('>>>> Testset: {}...'.format(testset))
data_loader_test = torch.utils.data.DataLoader(
dataset=MyData(testset, image_size=config.size, is_train=False),
batch_size=config.batch_size_valid, shuffle=False, num_workers=config.num_workers, pin_memory=True
)
for weights in weights_lst:
if int(weights.strip('.pth').split('epoch_')[-1]) % 1 != 0:
continue
print('\tInferencing {}...'.format(weights))
state_dict = torch.load(weights, map_location='cpu', weights_only=True)
state_dict = check_state_dict(state_dict)
model.load_state_dict(state_dict)
model = model.to(device)
inference(
model, data_loader_test=data_loader_test, pred_root=args.pred_root,
method='--'.join([w.rstrip('.pth') for w in weights.split(os.sep)[-2:]]),
testset=testset, device=config.device
)
if __name__ == '__main__':
# Parameter from command line
parser = argparse.ArgumentParser(description='')
parser.add_argument('--ckpt', type=str, help='model folder')
parser.add_argument('--ckpt_folder', default=sorted(glob(os.path.join('ckpt', '*')))[-1], type=str, help='model folder')
parser.add_argument('--pred_root', default='e_preds', type=str, help='Output folder')
parser.add_argument('--testsets',
default=config.testsets.replace(',', '+'),
type=str,
help="Test all sets: DIS5K -> 'DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'")
args = parser.parse_args()
if config.precisionHigh:
torch.set_float32_matmul_precision('high')
main(args)
|