Spaces:
Runtime error
Runtime error
File size: 4,741 Bytes
06e6b8c | 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 | import os
import argparse
from glob import glob
from tqdm import tqdm
import cv2
import torch
from contextlib import nullcontext
from dataset import MyData
from models.birefnet import BiRefNet
from utils import save_tensor_img, check_state_dict
from config import Config
config = Config()
mixed_precision = config.mixed_precision
if mixed_precision == 'fp16':
mixed_dtype = torch.float16
elif mixed_precision == 'bf16':
mixed_dtype = torch.bfloat16
else:
mixed_dtype = None
autocast_ctx = torch.amp.autocast(device_type='cuda', dtype=mixed_dtype) if mixed_dtype else nullcontext()
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 config.verbose_eval else data_loader_test:
inputs = batch[0].to(device)
label_paths = batch[-1]
with autocast_ctx, torch.no_grad():
scaled_preds = model(inputs)[-1].sigmoid().to(torch.float32)
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):
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)
else:
print('Undefined model: {}.'.format(config.model))
return None
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
)
try:
if args.resolution in [None, 'None', 0, '']:
# Use original resolution for inference.
data_size = None
elif args.resolution in ['config.size']:
data_size = config.size
else:
data_size = [int(l) for l in args.resolution.split('x')]
except Exception as e:
print(f"Exception: {type(e).__name__} at line {e.__traceback__.tb_lineno} of {__file__}: {e}")
# default as the config.size.
data_size = config.size
for testset in args.testsets.split('+'):
print('>>>> Testset: {}...'.format(testset))
data_loader_test = torch.utils.data.DataLoader(
dataset=MyData(testset, data_size=data_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:]]) + '-reso_{}'.format('x'.join([str(s) for s in data_size])),
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('ckpts', '*')))[-1], type=str, help='model folder')
parser.add_argument('--pred_root', default='e_preds', type=str, help='Output folder')
parser.add_argument('--resolution', default='default', type=str, help='WeixHei')
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)
|