def update_supervised(self, imL_src, imR_src, dispL):
assert len(imL_src.shape) == 4
bn, c, h0, w0 = imL_src.shape
sw, sh, ew, eh = mycrop(w0, h0, self.size_crop, size_cell=64)
# image prepare
imL = imL_src[:, :, sh:eh, sw:ew]
imR = imR_src[:, :, sh:eh, sw:ew]
dispL = dispL[:, :, sh:eh, sw:ew]
imL = to_tensor(imL, volatile=False, requires_grad=False)
imR = to_tensor(imR, volatile=False, requires_grad=False)
dispL = to_tensor(dispL, volatile=False, requires_grad=False)
# disp estimate
scale_dispLs, dispLs = self.predict(imL, imR)
# compute loss and backward
args = {
"disp_gt": dispL,
"disps": dispLs,
"scale_disps": scale_dispLs,
"flag_smooth": True,
}
loss = self.lossfun(args)
self.optim.zero_grad()
loss.backward()
self.optim.step()
self.updated += 1
self.lr_adjust()
return loss.data[0]
def test():
dirpath = r"/media/qjc/D/data/testimgs_stereo/"
# img1 = load_image(dirpath + "im0.png").transpose(2, 0, 1)[None]/255.0
# img2 = load_image(dirpath + "im1.png").transpose(2, 0, 1)[None]/255.0
# disp1 = load_disp(dirpath + "disp0.pfm")[None, None]/1.0
# disp2 = load_disp(dirpath + "disp1.pfm")[None, None]/1.0
img1 = load_image(dirpath + "im2.ppm").transpose(2, 0, 1)[None] / 255.0
img2 = load_image(dirpath + "im6.ppm").transpose(2, 0, 1)[None] / 255.0
disp1 = load_disp(dirpath + "disp2.pgm")[None, None] / 8.0
disp2 = load_disp(dirpath + "disp6.pgm")[None, None] / 8.0
# img1 = load_image(dirpath + "10L.png").transpose(2, 0, 1)[None]/255.0
# img2 = load_image(dirpath + "10R.png").transpose(2, 0, 1)[None]/255.0
# disp1 = load_disp(dirpath + "disp10L.png")[None, None]
# disp2 = load_disp(dirpath + "disp10R.png")[None, None]
im1 = to_tensor(img1)
im2 = to_tensor(img2)
d1 = to_tensor(disp1)
d2 = to_tensor(disp2)
im1t = imwrap_BCHW(im2, -d1)
im2t = imwrap_BCHW(im1, d2)
ssim = SSIM(window_size=11)
ssim1 = ssim(im1, im1t)
ssim2 = ssim(im2, im2t)
ssim3 = ssim(im1, im2)
abs1 = torch.abs(im1 - im1t).sum(dim=1, keepdim=True)
abs2 = torch.abs(im2 - im2t).sum(dim=1, keepdim=True)
print ssim1.shape, ssim2.shape
print ssim1.mean().data[0], ssim2.mean().data[0], ssim3.mean().data[0]
imsplot_tensor(im1, im2, im1t, im2t, 1 - ssim1, 1 - ssim2, abs1, abs2)
plt.show()
def update_self(self, imL_src, imR_src):
assert len(imL_src.shape) == 4
bn, c, h0, w0 = imL_src.shape
sw, sh, ew, eh = mycrop(w0, h0, self.size_crop, size_cell=64)
# image prepare
imL = imL_src[:, :, sh:eh, sw:ew]
imR = imR_src[:, :, sh:eh, sw:ew]
imL = to_tensor(imL, volatile=False, requires_grad=False)
imR = to_tensor(imR, volatile=False, requires_grad=False)
# disp estimate
dispLs = self.model(imL, imL)
dispRs = self.model(imR, imR)
# compute loss and backward
loss = 0
for disp in dispLs:
loss = loss + disp.mean()
for disp in dispRs:
loss = loss + disp.mean()
self.optim.zero_grad()
loss.backward()
self.optim.step()
def update_selfsupervised(self, imL_src, imR_src):
assert len(imL_src.shape) == 4
if (torch.is_tensor(imL_src)):
imL_src = imL_src.numpy()
imR_src = imR_src.numpy()
bn, c, h0, w0 = imL_src.shape
sw, sh, ew, eh = mycrop(w0, h0, self.size_crop, size_cell=64)
# image prepare
imL = imL_src[:, :, sh:eh, sw:ew]
imR = imR_src[:, :, sh:eh, sw:ew]
imL1 = np.flip(imR, axis=-1).copy()
imR1 = np.flip(imL, axis=-1).copy()
imL = to_tensor(imL, volatile=False, requires_grad=False)
imR = to_tensor(imR, volatile=False, requires_grad=False)
imL1 = to_tensor(imL1, volatile=False, requires_grad=False)
imR1 = to_tensor(imR1, volatile=False, requires_grad=False)
# disp estimate
scale_dispLs, dispLs = self.predict(imL, imR)
scale_dispL1s, dispL1s = self.predict(imL1, imR1)
# compute loss and backward
if (self.lossfun.flag_mask):
imR_src = to_tensor(imR_src, requires_grad=False)
imR1_src = to_tensor(np.flip(imL_src, axis=-1).copy(),
requires_grad=False)
LeftTop = [sw, sh]
LeftTop1 = [w0 - ew, sh]
else:
imR_src = imR
imR1_src = imR1
LeftTop = [0, 0]
LeftTop1 = [0, 0]
args = {
"imR_src": imR_src,
"imL": imL,
"dispLs": dispLs,
"scale_dispLs": scale_dispLs,
"LeftTop": LeftTop,
"imR1_src": imR1_src,
"imL1": imL1,
"dispL1s": dispL1s,
"scale_dispL1s": scale_dispL1s,
"LeftTop1": LeftTop1,
}
loss = self.lossfun(args)
assert loss.data[0] > 0
self.optim.zero_grad()
loss.backward()
self.optim.step()
self.updated += 1
self.lr_adjust()
return loss.data[0]
def augment_color_pair(self, left_image, right_image):
# randomly augment image
bn, c = left_image.shape[:2]
# randomly shift gamma
do_augment = torch.rand(1)[0]
if (do_augment > 0.5):
random_gamma = to_tensor(torch.rand(bn) * 0.4 + 0.8)
left_image = (left_image.transpose(0, -1)**random_gamma).transpose(
0, -1)
random_gamma = to_tensor(torch.rand(bn) * 0.4 + 0.8)
right_image = (right_image.transpose(0,
-1)**random_gamma).transpose(
0, -1)
# randomly shift brightness
do_augment = torch.rand(1)[0]
if (do_augment > 0.5):
random_brightness = to_tensor(torch.rand(bn) * 1.5 + 0.5)
left_image = (left_image.transpose(0, -1) *
random_brightness).transpose(0, -1)
random_brightness = to_tensor(torch.rand(bn) * 1.5 + 0.5)
right_image = (right_image.transpose(0, -1) *
random_brightness).transpose(0, -1)
# randomly shift color
do_augment = torch.rand(1)[0]
if (do_augment > 0.5):
random_colors = torch.rand(bn, c) * 0.4 + 0.8
color_image = (torch.ones(left_image.shape).permute(2, 3, 0, 1) *
random_colors).permute(2, 3, 0, 1)
left_image = left_image * to_tensor(color_image)
random_colors = torch.rand(bn, c) * 0.4 + 0.8
color_image = (torch.ones(left_image.shape).permute(2, 3, 0, 1) *
random_colors).permute(2, 3, 0, 1)
right_image = right_image * to_tensor(color_image)
# saturate
do_augment = torch.rand(1)[0]
if (do_augment > 0.5):
left_image = left_image.clamp(0, 1)
right_image = right_image.clamp(0, 1)
return left_image, right_image
def estimate(self, imL, imR):
imL = to_tensor(imL, volatile=True, requires_grad=False)
imR = to_tensor(imR, volatile=True, requires_grad=False)
_, dispLs = self.model(imL, imR)
return to_numpy(dispLs[0])
import matplotlib.pyplot as plt
from util.util import to_tensor, imsplot_tensor
from dataloader.imagerw import load_image, load_disp #, save_image
from loss import losses
if __name__ == "__main__":
dirpath = r"/media/qjc/D/data/testimgs_stereo/"
# img1 = load_image(dirpath + "im0.png").transpose(2, 0, 1)[None]/255.0
# img2 = load_image(dirpath + "im1.png").transpose(2, 0, 1)[None]/255.0
# disp1 = load_disp(dirpath + "disp0.pfm")[None, None]/1.0
# disp2 = load_disp(dirpath + "disp1.pfm")[None, None]/1.0
img1 = load_image(dirpath + "im2.ppm").transpose(2, 0, 1)[None]/255.0
img2 = load_image(dirpath + "im6.ppm").transpose(2, 0, 1)[None]/255.0
disp1 = load_disp(dirpath + "disp2.pgm")[None, None]/8.0
disp2 = load_disp(dirpath + "disp6.pgm")[None, None]/8.0
img1t = to_tensor(np.flip(img2, axis=-1).copy())
img2t = to_tensor(np.flip(img1, axis=-1).copy())
img1 = to_tensor(img1, requires_grad=True)
img2 = to_tensor(img2, requires_grad=True)
disp1 = to_tensor(disp1, requires_grad=True)
disp1t = to_tensor(np.flip(disp2, axis=-1).copy(), requires_grad=True)
lossfun = losses(loss_name="depthmono")
"imR_src, imL, dispLs, scale_dispLs, LeftTop, imR1_src, imL1, dispL1s, scale_dispL1s, LeftTop1"
args = {
"imR_src": img2, "imL": img1, "dispLs": [disp1], "scale_dispLs": [0], "LeftTop": [0, 0],
"imR1_src": img2t, "imL1": img1t, "dispL1s": [disp1t], "scale_dispL1s": [0], "LeftTop1": [0, 0],
}
loss = lossfun(args)
loss.backward(retain_graph=True)
print loss.data[0]