from data_loader import CenterCrop
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset
from model import BASNet
import pytorch_ssim
import pytorch_iou
from utils import AverageMeter, VisdomLinePlotter
from paperBASNET.metrics import getMAE, getPRCurve, getMaxFMeasure, getRelaxedFMeasure, own_RelaxedFMeasure
# ------- 1. define loss function --------
bce_loss = nn.BCELoss(size_average=True)
ssim_loss = pytorch_ssim.SSIM(window_size=11, size_average=True)
iou_loss = pytorch_iou.IOU(size_average=True)
def normPRED(d):
ma = torch.max(d)
mi = torch.min(d)
dn = (d - mi) / (ma - mi)
return dn
def bce_ssim_loss(pred, target):
bce_out = bce_loss(pred, target)
ssim_out = 1 - ssim_loss(pred, target)
image_root = './train_dataset/image/'
depth_root = './train_dataset/depth/'
gt_root = './train_dataset/GT/'
# image_root = './train_dataset_DMRA/image/'
# depth_root = './train_dataset_DMRA/depth/'
# gt_root = './train_dataset_DMRA/GT/'
train_loader = get_loader(image_root,
depth_root,
gt_root,
batchsize=opt.batchsize,
trainsize=opt.trainsize)
total_step = len(train_loader)
CE = torch.nn.BCEWithLogitsLoss()
IOU = pytorch_iou.IOU(size_average=True)
def train(train_loader, model, optimizer, epoch):
model.train()
for i, pack in enumerate(train_loader, start=1):
optimizer.zero_grad()
images, depths, gts = pack
images = Variable(images)
depths = Variable(depths)
gts = Variable(gts)
images = images.cuda()
depths = depths.cuda()
gts = gts.cuda()
s1, s2, s3, s4, s5, s1_sig, s2_sig, s3_sig, s4_sig, s5_sig = model(
