def eval_F1(self):
# Reset f1 calc class
self.f1_class = F1Score(self.device)
for batch in self.eval_loader:
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
parts_mask_gt = batch['parts_mask_gt'].to(self.device)
orig, orig_label = batch['orig'].to(
self.device), batch['orig_label']
N, L, H, W = orig_label.shape
stage1_pred = F.softmax(self.model(image), dim=1)
assert stage1_pred.shape == (N, 9, 128, 128)
# Imshow stage1_pred on Tensorborad
stage1_pred_grid = torchvision.utils.make_grid(
stage1_pred.argmax(dim=1, keepdim=True))
self.writer.add_image("stage1 predict%s" % uuid,
stage1_pred_grid,
self.step,
dataformats="HW")
# Mask2Theta using ModelB
theta = self.select_net(stage1_pred)
assert theta.shape == (N, 6, 2, 3)
# AffineCrop
parts, parts_label, _ = affine_crop(img=orig,
label=orig_label,
theta_in=theta,
map_location=self.device)
assert parts.shape == (N, 6, 3, 81, 81)
# Predict Cropped Parts
stage2_pred = self.model2(parts)
# Calc stage2 CrossEntropy Loss
error = []
for i in range(6):
error.append(
self.criterion(stage2_pred[i], parts_mask_gt[:, i].long()))
# Imshow final predict mask
final_pred = affine_mapback(stage2_pred, theta, self.device)
final_grid = torchvision.utils.make_grid(
final_pred.argmax(dim=1, keepdim=True))
self.writer.add_image("final predict_%s" % uuid,
final_grid[0],
global_step=self.step,
dataformats='HW')
# Accumulate F1
self.f1_class.forward(final_pred,
orig_label.argmax(dim=1, keepdim=False))
# Calc F1 overall
_, F1_overall = self.f1_class.get_f1_score()
# clean f1_class
del self.f1_class
return F1_overall, error
def eval_error(self):
loss_list = []
for batch in self.eval_loader:
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
orig, orig_label = batch['orig'].to(
self.device), batch['orig_label']
# parts_mask_gt = batch['parts_mask_gt'].to(self.device)
N, L, H, W = orig_label.shape
stage1_pred = self.model(image)
assert stage1_pred.shape == (N, 9, 128, 128)
theta = self.select_net(F.softmax(stage1_pred, dim=1))
assert theta.shape == (N, 6, 2, 3)
parts, parts_label, _ = affine_crop(img=orig,
label=orig_label,
theta_in=theta,
map_location=self.device)
assert parts.shape == (N, 6, 3, 81, 81)
stage2_pred = self.model2(parts)
assert len(stage2_pred) == 6
loss = []
for i in range(6):
loss.append(
self.criterion(stage2_pred[i], parts_label[i].argmax(
dim=1, keepdim=False)).item())
# loss.append(self.criterion(stage2_pred[i], parts_mask_gt[:, i].long()).item())
loss_list.append(np.mean(loss))
return np.mean(loss_list)
def eval_error(self):
loss_list = []
for batch in self.eval_loader:
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
orig, orig_label = batch['orig'].to(
self.device), batch['orig_label'].to(self.device)
parts_mask = batch['parts_mask_gt'].to(self.device)
N, L, H, W = orig_label.shape
assert label.shape == (N, 9, 128, 128)
theta = self.select_net(label)
assert theta.shape == (N, 6, 2, 3)
parts, parts_label, _ = affine_crop(img=orig,
label=orig_label,
theta_in=theta,
map_location=self.device)
assert parts.grad_fn is not None
assert parts.shape == (N, 6, 3, 81, 81)
stage2_pred = self.model2(parts)
assert len(stage2_pred) == 6
loss = []
for i in range(6):
loss.append(
self.criterion(stage2_pred[i],
parts_mask[:, i].long()).item())
loss_list.append(np.mean(loss))
return np.mean(loss_list)
def train_loss(self, batch):
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
orig, orig_label = batch['orig'].to(
self.device), batch['orig_label'].to(self.device)
parts_mask = batch['parts_mask_gt'].to(self.device)
N, L, H, W = orig_label.shape
assert label.shape == (N, 9, 128, 128)
theta = self.select_net(label)
assert theta.shape == (N, 6, 2, 3)
parts, parts_label, _ = affine_crop(img=orig,
label=orig_label,
theta_in=theta,
map_location=self.device)
assert parts.grad_fn is not None
assert parts.shape == (N, 6, 3, 81, 81)
stage2_pred = self.model2(parts)
assert len(stage2_pred) == 6
loss = []
for i in range(6):
loss.append(self.criterion(stage2_pred[i], parts_mask[:,
i].long()))
loss = torch.stack(loss)
return loss
def train_loss(self, batch):
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
orig, orig_label = batch['orig'].to(self.device), batch['orig_label']
parts_mask_gt = batch['parts_mask_gt'].to(self.device)
N, L, H, W = orig_label.shape
assert parts_mask_gt.shape == (N, 6, 81, 81)
# Get Stage1 coarse mask
stage1_pred = F.softmax(self.model(image), dim=1)
assert stage1_pred.shape == (N, 9, 128, 128)
# Mask2Theta
theta = self.select_net(stage1_pred)
assert theta.shape == (N, 6, 2, 3)
# Affine_cropper
parts, parts_label, _ = affine_crop(img=orig,
label=orig_label,
theta_in=theta,
map_location=self.device)
assert parts.shape == (N, 6, 3, 81, 81)
# Make sure the backward grad stream unblocked
assert parts.grad_fn is not None
# Get Stage2 Mask Predict
stage2_pred = self.model2(parts)
# Calc Stage2 CrossEntropy Loss
loss = []
for i in range(6):
loss.append(
self.criterion(stage2_pred[i],
parts_label[i].argmax(dim=1, keepdim=False)))
# loss.append(self.criterion(stage2_pred[i], parts_mask_gt[:, i].long()))
loss = torch.stack(loss)
return loss
def eval_F1(self):
step = 0
f1_class = F1Score(self.device) # reload f1 calc class
loss_list = []
for batch in self.eval_loader:
step += 1
image, label = batch['image'].to(self.device), batch['labels'].to(
self.device)
orig, orig_label = batch['orig'].to(
self.device), batch['orig_label'].to(self.device)
N, L, H, W = orig_label.shape
# Get stage1 predict mask (corase mask)
stage1_pred = self.model(image)
assert stage1_pred.shape == (N, 9, 128, 128)
# Mask2Theta
theta = self.select_net(F.softmax(stage1_pred, dim=1))
test_pred = F.softmax(stage1_pred, dim=1).argmax(dim=1,
keepdim=True)
test_pred_grid = torchvision.utils.make_grid(test_pred)
self.writer.add_image("stage1_pred_%s" % uuid,
test_pred_grid[0],
step,
dataformats='HW')
assert theta.shape == (N, 6, 2, 3)
"""""
Using original mask groundtruth to calc theta_groundtruth
""" ""
assert orig_label.shape == (N, 9, 1024, 1024)
cens = torch.floor(calc_centroid(orig_label))
assert cens.shape == (N, 9, 2)
points = torch.floor(
torch.cat(
[cens[:, 1:6], cens[:, 6:9].mean(dim=1, keepdim=True)],
dim=1))
theta_label = torch.zeros((N, 6, 2, 3),
device=self.device,
requires_grad=False)
for i in range(6):
theta_label[:, i, 0, 0] = (81. - 1.) / (W - 1)
theta_label[:, i, 0,
2] = -1. + (2. * points[:, i, 1]) / (W - 1)
theta_label[:, i, 1, 1] = (81. - 1.) / (H - 1)
theta_label[:, i, 1,
2] = -1. + (2. * points[:, i, 0]) / (H - 1)
"""""
Calc Regression loss, Loss func: Smooth L1 loss
""" ""
loss = self.regress_loss(theta, theta_label)
loss_list.append(loss.item())
"""""
Imshow parts cropped by Affine cropper(using predicted theta) for debug convinience
""" ""
temp = []
for i in range(theta.shape[1]):
test = theta[:, i]
grid = F.affine_grid(theta=test,
size=[N, 3, 81, 81],
align_corners=True)
temp.append(
F.grid_sample(input=orig, grid=grid, align_corners=True))
parts = torch.stack(temp, dim=1)
assert parts.shape == (N, 6, 3, 81, 81)
for i in range(6):
parts_grid = torchvision.utils.make_grid(
parts[:, i].detach().cpu())
self.writer.add_image('croped_parts_%s_%d' % (uuid, i),
parts_grid, step)
"""""
Imshow final predict mask for debug convinience
""" ""
parts, parts_labels, _ = affine_crop(orig,
orig_label,
theta_in=theta,
map_location=self.device)
final_parts_mask = affine_mapback(parts_labels, theta, self.device)
final_grid = torchvision.utils.make_grid(
final_parts_mask.argmax(dim=1, keepdim=True))
self.writer.add_image("final_parts_mask %s" % uuid,
final_grid[0],
global_step=step,
dataformats='HW')
# Accumulate F1
f1_class.forward(final_parts_mask,
orig_label.argmax(dim=1, keepdim=False))
# Calc F1 overall
_, F1_overall = f1_class.get_f1_score()
return F1_overall, np.mean(loss_list)
for batch in dataloader['test']:
step += 1
orig = batch['orig'].to(device)
orig_label = batch['orig_label'].to(device)
image = batch['image'].to(device)
label = batch['labels'].to(device)
N, L, H, W = orig_label.shape
stage1_pred = model1(image)
assert stage1_pred.shape == (N, 9, 128, 128)
theta = select_model(F.softmax(stage1_pred, dim=1))
# cens = calc_centroid(orig_label)
# assert cens.shape == (N, 9, 2)
parts, parts_labels, _ = affine_crop(orig,
orig_label,
theta_in=theta,
map_location=device)
stage2_pred = model2(parts)
for i in range(6):
parts_grid = torchvision.utils.make_grid(parts[:, i].detach().cpu())
writer.add_image('croped_parts_%s_%d' % ("testABC", i), parts_grid,
step)
for i in range(6):
pred_grid = torchvision.utils.make_grid(stage2_pred[i].argmax(
dim=1, keepdim=True))
writer.add_image('ABC_stage2 predict_%d' % i,
pred_grid[0],
global_step=step,
dataformats='HW')
dataloader = {
x: DataLoader(Dataset[x], batch_size=4, shuffle=True, num_workers=4)
for x in ['train', 'val', 'test']
}
f1_class = F1Score(device)
# show predicts
step = 0
for batch in dataloader['test']:
step += 1
orig = batch['orig'].to(device)
orig_label = batch['orig_label'].to(device)
N, L, H, W = orig_label.shape
assert orig_label.shape == (N, 9, 1024, 1024)
cens = calc_centroid(orig_label)
assert cens.shape == (N, 9, 2)
parts, parts_labels, theta = affine_crop(orig, orig_label, cens, device)
stage2_pred = model2(parts)
for i in range(6):
pred_grid = torchvision.utils.make_grid(stage2_pred[i].argmax(
dim=1, keepdim=True))
writer.add_image('stage2 predict_%d' % i,
pred_grid[0],
global_step=step,
dataformats='HW')
final_pred = affine_mapback(stage2_pred, theta, device)
final_grid = torchvision.utils.make_grid(
final_pred.argmax(dim=1, keepdim=True))
writer.add_image("final predict",
final_grid[0],
step += 1
orig = batch['orig'].to(device)
orig_label = batch['orig_label'].to(device)
image = batch['image'].to(device)
label = batch['labels'].to(device)
parts_gt = batch['parts_gt'].to(device)
parts_mask_gt = batch['parts_mask_gt'].to(device)
N, L, H, W = orig_label.shape
cens = torch.floor(calc_centroid(orig_label))
# Test B
# theta = select_model(label)
parts, parts_labels, theta = affine_crop(orig,
orig_label,
points=cens,
map_location=device)
# parts, parts_labels, _ = affine_crop(orig, orig_label, theta_in=theta, map_location=device)
stage2_pred = model2(parts_gt)
# imshow predict
for i in range(6):
pred_grid = torchvision.utils.make_grid(stage2_pred[i].argmax(
dim=1, keepdim=True))
writer.add_image('stage2 predict_%d' % i,
pred_grid[0],
global_step=step,
dataformats='HW')
# imshow crop