help='every N epochs save your trained snapshot')
parser.add_argument('--save_model',
type=str,
default='./Snapshot/2020-CVPR-SINet/')
parser.add_argument('--train_img_dir',
type=str,
default='/content/drive/MyDrive/Test12/input/')
parser.add_argument('--train_gt_dir',
type=str,
default='/content/drive/MyDrive/Test12/groundtruth/')
opt = parser.parse_args()
torch.cuda.set_device(0)
# TIPS: you also can use deeper network for better performance like channel=64
model_SINet = SINet_ResNet50(channel=32).cuda()
print('-' * 30, model_SINet, '-' * 30)
optimizer = torch.optim.Adam(model_SINet.parameters(), opt.lr)
LogitsBCE = torch.nn.BCEWithLogitsLoss()
net, optimizer = amp.initialize(model_SINet, optimizer,
opt_level='O1') # NOTES: Ox not 0x
train_loader = get_loader(opt.train_img_dir,
opt.train_gt_dir,
batchsize=opt.batchsize,
trainsize=opt.trainsize,
num_workers=12)
total_step = len(train_loader)
import argparse
from Src.SINet import SINet_ResNet50
from Src.utils.Dataloader import test_dataset
import cv2 as cv
import time
parser = argparse.ArgumentParser()
parser.add_argument('--testsize', type=int, default=352, help='the snapshot input size')
parser.add_argument('--model_path', type=str,
default='./Snapshot/2020-CVPR-SINet/SINet_40.pth')
parser.add_argument('--test_save', type=str,
default='./Result/2020-CVPR-SINet-New/')
parser.add_argument('--save_all', type=bool, default=True)
opt = parser.parse_args()
model = SINet_ResNet50().cpu()
model.load_state_dict(torch.load(opt.model_path, map_location=torch.device('cpu')))
model.eval()
close = False
for dataset in ['COD10K-v3', 'MYTEST', 'CAMO', 'CHAMELEON', 'COD10K']:
if close:
break
save_path = opt.test_save + dataset + '/'
os.makedirs(save_path, exist_ok=True)
imgpath = './Dataset/TestDataset/{}/Imgs/'.format(dataset)
gtpath = './Dataset/TestDataset/{}/GT/'.format(dataset)
if dataset == 'COD10K-v3':
imgpath = './Dataset/TestDataset/{}/Image/'.format(dataset)
gtpath = './Dataset/TestDataset/{}/GT_Instance/'.format(dataset)
from Src.utils.trainer import eval_mae, numpy2tensor
parser = argparse.ArgumentParser()
parser.add_argument('--testsize',
type=int,
default=352,
help='the snapshot input size')
parser.add_argument('--model_path',
type=str,
default='./Snapshot/2020-CVPR-SINet/SINet_40.pth')
parser.add_argument('--test_save',
type=str,
default='./Result/2020-CVPR-SINet-New/')
opt = parser.parse_args()
model = SINet_ResNet50().cuda()
model.load_state_dict(torch.load(opt.model_path))
model.eval()
for dataset in ['COD10K']:
save_path = opt.test_save + dataset + '/'
os.makedirs(save_path, exist_ok=True)
# NOTES:
# if you plan to inference on your customized dataset without grouth-truth,
# you just modify the params (i.e., `image_root=your_test_img_path` and `gt_root=your_test_img_path`)
# with the same filepath. We recover the original size according to the shape of grouth-truth, and thus,
# the grouth-truth map is unnecessary actually.
test_loader = test_dataset(
image_root='./Dataset/TestDataset/{}/Image/'.format(dataset),
gt_root='./Dataset/TestDataset/{}/GT/'.format(dataset),
testsize=opt.testsize)