bn, SGNum, _, envRow, envCol = axisPred.size() envmapsPred = torch.cat([ axisPred.view(bn, SGNum * 3, envRow, envCol), lambPred, weightPred ], dim=1) envmapsPreds.append(envmapsPred) envmapsPredImage, axisPred, lambPred, weightPred = output2env.output2env( axisPred, lambPred, weightPred) envmapsPredImages.append(envmapsPredImage) diffusePred, specularPred = renderLayer.forwardEnv( albedoPreds[0], normalPreds[0], roughPreds[0], envmapsPredImages[0]) diffusePredNew, specularPredNew = models.LSregressDiffSpec( diffusePred, specularPred, imBatchSmall, diffusePred, specularPred) renderedPred = diffusePredNew + specularPredNew renderedPreds.append(renderedPred) cDiff, cSpec = (torch.sum(diffusePredNew) / torch.sum(diffusePred)).data.item(), ( (torch.sum(specularPredNew)) / (torch.sum(specularPred))).data.item() if cSpec < 1e-3: cAlbedo = 1 / albedoPreds[-1].max().data.item() cLight = cDiff / cAlbedo else: cLight = cSpec cAlbedo = cDiff / cLight cAlbedo = np.clip(cAlbedo, 1e-3, 1 / albedoPreds[-1].max().data.item())
mode='bilinear') if roughPreBatch.size(2) < opt.imHeight or roughPreBatch.size( 3) < opt.imWidth: roughPreBatch = F.interpolate(roughPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if depthPreBatch.size(2) < opt.imHeight or depthPreBatch.size( 3) < opt.imWidth: depthPreBatch = F.interpolate(depthPreBatch, [opt.imHeight, opt.imWidth]) # Regress the diffusePred and specular Pred envRow, envCol = diffusePreBatch.size(2), diffusePreBatch.size(3) imBatchSmall = F.adaptive_avg_pool2d(imBatch, (envRow, envCol)) diffusePreBatch, specularPreBatch = models.LSregressDiffSpec( diffusePreBatch, specularPreBatch, imBatchSmall, diffusePreBatch, specularPreBatch) if diffusePreBatch.size(2) < opt.imHeight or diffusePreBatch.size( 3) < opt.imWidth: diffusePreBatch = F.interpolate(diffusePreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if specularPreBatch.size(2) < opt.imHeight or specularPreBatch.size( 3) < opt.imWidth: specularPreBatch = F.interpolate(specularPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') renderedImBatch = diffusePreBatch + specularPreBatch
def wrapperBRDFLight(dataBatch, opt, encoder, albedoDecoder, normalDecoder, roughDecoder, depthDecoder, lightEncoder, axisDecoder, lambDecoder, weightDecoder, output2env, renderLayer, offset=1.0, isLightOut=False): # Load data from cpu to gpu albedo_cpu = dataBatch['albedo'] albedoBatch = Variable(albedo_cpu).cuda() normal_cpu = dataBatch['normal'] normalBatch = Variable(normal_cpu).cuda() rough_cpu = dataBatch['rough'] roughBatch = Variable(rough_cpu).cuda() depth_cpu = dataBatch['depth'] depthBatch = Variable(depth_cpu).cuda() segArea_cpu = dataBatch['segArea'] segEnv_cpu = dataBatch['segEnv'] segObj_cpu = dataBatch['segObj'] seg_cpu = torch.cat([segArea_cpu, segEnv_cpu, segObj_cpu], dim=1) segBatch = Variable(seg_cpu).cuda() segBRDFBatch = segBatch[:, 2:3, :, :] segAllBatch = segBatch[:, 0:1, :, :] + segBatch[:, 2:3, :, :] # Load the image from cpu to gpu im_cpu = (dataBatch['im']) imBatch = Variable(im_cpu).cuda() if opt.cascadeLevel > 0: albedoPre_cpu = dataBatch['albedoPre'] albedoPreBatch = Variable(albedoPre_cpu).cuda() normalPre_cpu = dataBatch['normalPre'] normalPreBatch = Variable(normalPre_cpu).cuda() roughPre_cpu = dataBatch['roughPre'] roughPreBatch = Variable(roughPre_cpu).cuda() depthPre_cpu = dataBatch['depthPre'] depthPreBatch = Variable(depthPre_cpu).cuda() diffusePre_cpu = dataBatch['diffusePre'] diffusePreBatch = Variable(diffusePre_cpu).cuda() specularPre_cpu = dataBatch['specularPre'] specularPreBatch = Variable(specularPre_cpu).cuda() if albedoPreBatch.size(2) < opt.imHeight or albedoPreBatch.size( 3) < opt.imWidth: albedoPreBatch = F.interpolate(albedoPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if normalPreBatch.size(2) < opt.imHeight or normalPreBatch.size( 3) < opt.imWidth: normalPreBatch = F.interpolate(normalPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if roughPreBatch.size(2) < opt.imHeight or roughPreBatch.size( 3) < opt.imWidth: roughPreBatch = F.interpolate(roughPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if depthPreBatch.size(2) < opt.imHeight or depthPreBatch.size( 3) < opt.imWidth: depthPreBatch = F.interpolate(depthPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') # Regress the diffusePred and specular Pred envRow, envCol = diffusePreBatch.size(2), diffusePreBatch.size(3) imBatchSmall = F.adaptive_avg_pool2d(imBatch, (envRow, envCol)) diffusePreBatch, specularPreBatch = models.LSregressDiffSpec( diffusePreBatch.detach(), specularPreBatch.detach(), imBatchSmall, diffusePreBatch, specularPreBatch) if diffusePreBatch.size(2) < opt.imHeight or diffusePreBatch.size( 3) < opt.imWidth: diffusePreBatch = F.interpolate(diffusePreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if specularPreBatch.size(2) < opt.imHeight or specularPreBatch.size( 3) < opt.imWidth: specularPreBatch = F.interpolate(specularPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') envmapsPre_cpu = dataBatch['envmapsPre'] envmapsPreBatch = Variable(envmapsPre_cpu).cuda() # Normalize Albedo and depth bn, ch, nrow, ncol = albedoPreBatch.size() albedoPreBatch = albedoPreBatch.view(bn, -1) albedoPreBatch = albedoPreBatch / torch.clamp( torch.mean(albedoPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 albedoPreBatch = albedoPreBatch.view(bn, ch, nrow, ncol) bn, ch, nrow, ncol = depthPreBatch.size() depthPreBatch = depthPreBatch.view(bn, -1) depthPreBatch = depthPreBatch / torch.clamp( torch.mean(depthPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 depthPreBatch = depthPreBatch.view(bn, ch, nrow, ncol) envmaps_cpu = dataBatch['envmaps'] envmapsBatch = Variable(envmaps_cpu).cuda() envmapsInd_cpu = dataBatch['envmapsInd'] envmapsIndBatch = Variable(envmapsInd_cpu).cuda() ######################################################## # Build the cascade network architecture # if opt.cascadeLevel == 0: inputBatch = imBatch elif opt.cascadeLevel > 0: inputBatch = torch.cat([ imBatch, albedoPreBatch, normalPreBatch, roughPreBatch, depthPreBatch, diffusePreBatch, specularPreBatch ], dim=1) # Initial Prediction x1, x2, x3, x4, x5, x6 = encoder(inputBatch) albedoPred = 0.5 * (albedoDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) normalPred = normalDecoder(imBatch, x1, x2, x3, x4, x5, x6) roughPred = roughDecoder(imBatch, x1, x2, x3, x4, x5, x6) depthPred = 0.5 * (depthDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) albedoBatch = segBRDFBatch * albedoBatch albedoPred1 = models.LSregress( albedoPred.detach() * segBRDFBatch.expand_as(albedoPred), albedoBatch * segBRDFBatch.expand_as(albedoBatch), albedoPred) albedoPred1 = torch.clamp(albedoPred1, 0, 1) depthPred1 = models.LSregress( depthPred.detach() * segAllBatch.expand_as(depthPred), depthBatch * segAllBatch.expand_as(depthBatch), depthPred) ## Compute Errors pixelObjNum = (torch.sum(segBRDFBatch).cpu().data).item() pixelAllNum = (torch.sum(segAllBatch).cpu().data).item() albedoErr = torch.sum( (albedoPred1 - albedoBatch) * (albedoPred1 - albedoBatch) * segBRDFBatch.expand_as(albedoBatch) / pixelObjNum / 3.0) normalErr = torch.sum( (normalPred - normalBatch) * (normalPred - normalBatch) * segAllBatch.expand_as(normalBatch)) / pixelAllNum / 3.0 roughErr = torch.sum((roughPred - roughBatch) * (roughPred - roughBatch) * segBRDFBatch) / pixelObjNum depthErr = torch.sum( (torch.log(depthPred1 + 1) - torch.log(depthBatch + 1)) * (torch.log(depthPred1 + 1) - torch.log(depthBatch + 1)) * segAllBatch.expand_as(depthBatch)) / pixelAllNum # Normalize Albedo and depth bn, ch, nrow, ncol = albedoPred.size() albedoPred = albedoPred.view(bn, -1) albedoPred = albedoPred / torch.clamp(torch.mean(albedoPred, dim=1), min=1e-10).unsqueeze(1) / 3.0 albedoPred = albedoPred.view(bn, ch, nrow, ncol) bn, ch, nrow, ncol = depthPred.size() depthPred = depthPred.view(bn, -1) depthPred = depthPred / torch.clamp(torch.mean(depthPred, dim=1), min=1e-10).unsqueeze(1) / 3.0 depthPred = depthPred.view(bn, ch, nrow, ncol) imBatchLarge = F.interpolate(imBatch, [480, 640], mode='bilinear') albedoPredLarge = F.interpolate(albedoPred, [480, 640], mode='bilinear') normalPredLarge = F.interpolate(normalPred, [480, 640], mode='bilinear') roughPredLarge = F.interpolate(roughPred, [480, 640], mode='bilinear') depthPredLarge = F.interpolate(depthPred, [480, 640], mode='bilinear') inputBatch = torch.cat([ imBatchLarge, albedoPredLarge, 0.5 * (normalPredLarge + 1), 0.5 * (roughPredLarge + 1), depthPredLarge ], dim=1) if opt.cascadeLevel == 0: x1, x2, x3, x4, x5, x6 = lightEncoder(inputBatch.detach()) elif opt.cascadeLevel > 0: x1, x2, x3, x4, x5, x6 = lightEncoder(inputBatch.detach(), envmapsPreBatch.detach()) # Prediction axisPred = axisDecoder(x1, x2, x3, x4, x5, x6, envmapsBatch) lambPred = lambDecoder(x1, x2, x3, x4, x5, x6, envmapsBatch) weightPred = weightDecoder(x1, x2, x3, x4, x5, x6, envmapsBatch) bn, SGNum, _, envRow, envCol = axisPred.size() envmapsPred = torch.cat( [axisPred.view(bn, SGNum * 3, envRow, envCol), lambPred, weightPred], dim=1) imBatchSmall = F.adaptive_avg_pool2d(imBatch, (opt.envRow, opt.envCol)) segBatchSmall = F.adaptive_avg_pool2d(segBRDFBatch, (opt.envRow, opt.envCol)) notDarkEnv = (torch.mean(torch.mean(torch.mean(envmapsBatch, 4), 4), 1, True) > 0.001).float() segEnvBatch = ( segBatchSmall * envmapsIndBatch.expand_as(segBatchSmall)).unsqueeze(-1).unsqueeze(-1) segEnvBatch = segEnvBatch * notDarkEnv.unsqueeze(-1).unsqueeze(-1) # Compute the recontructed error envmapsPredImage, axisPred, lambPred, weightPred = output2env.output2env( axisPred, lambPred, weightPred) pixelNum = max((torch.sum(segEnvBatch).cpu().data).item(), 1e-5) envmapsPredScaledImage = models.LSregress( envmapsPredImage.detach() * segEnvBatch.expand_as(envmapsBatch), envmapsBatch * segEnvBatch.expand_as(envmapsBatch), envmapsPredImage) reconstErr = torch.sum( ( torch.log(envmapsPredScaledImage + offset ) - torch.log(envmapsBatch + offset ) ) * ( torch.log(envmapsPredScaledImage + offset ) - torch.log(envmapsBatch + offset ) ) * segEnvBatch.expand_as(envmapsPredImage ) ) \ / pixelNum / 3.0 / opt.envWidth / opt.envHeight # Compute the rendered error pixelNum = max((torch.sum(segBatchSmall).cpu().data).item(), 1e-5) diffusePred, specularPred = renderLayer.forwardEnv(albedoPred.detach(), normalPred, roughPred, envmapsPredImage) diffusePredScaled, specularPredScaled = models.LSregressDiffSpec( diffusePred.detach(), specularPred.detach(), imBatchSmall, diffusePred, specularPred) renderedImPred = torch.clamp(diffusePredScaled + specularPredScaled, 0, 1) renderErr = torch.sum( (renderedImPred - imBatchSmall) * (renderedImPred - imBatchSmall) * segBatchSmall.expand_as(imBatchSmall ) ) \ / pixelNum / 3.0 if isLightOut == False: return [albedoPred, albedoErr, albedoBatch], \ [normalPred, normalErr, normalBatch ], \ [roughPred, roughErr, roughBatch], \ [depthPred, depthErr, depthBatch], \ [envmapsPredScaledImage, reconstErr, envmapsBatch], \ [renderedImPred, renderErr, imBatch] else: return [albedoPred, albedoErr, albedoBatch], \ [normalPred, normalErr, normalBatch], \ [roughPred, roughErr, roughBatch ], \ [depthPred, depthErr, depthBatch], \ [envmapsPredScaledImage, reconstErr], \ [renderedImPred, renderErr, imBatch], \ [envmapsPred, diffusePred, specularPred]
def wrapperBRDF(dataBatch, opt, encoder, albedoDecoder, normalDecoder, roughDecoder, depthDecoder): # Load data from cpu to gpu albedo_cpu = dataBatch['albedo'] albedoBatch = Variable(albedo_cpu).cuda() normal_cpu = dataBatch['normal'] normalBatch = Variable(normal_cpu).cuda() rough_cpu = dataBatch['rough'] roughBatch = Variable(rough_cpu).cuda() depth_cpu = dataBatch['depth'] depthBatch = Variable(depth_cpu).cuda() segArea_cpu = dataBatch['segArea'] segEnv_cpu = dataBatch['segEnv'] segObj_cpu = dataBatch['segObj'] seg_cpu = torch.cat([segArea_cpu, segEnv_cpu, segObj_cpu], dim=1) segBatch = Variable(seg_cpu).cuda() segBRDFBatch = segBatch[:, 2:3, :, :] segAllBatch = segBatch[:, 0:1, :, :] + segBatch[:, 2:3, :, :] # Load the image from cpu to gpu im_cpu = (dataBatch['im']) imBatch = Variable(im_cpu).cuda() if opt.cascadeLevel > 0: albedoPre_cpu = dataBatch['albedoPre'] albedoPreBatch = Variable(albedoPre_cpu).cuda() normalPre_cpu = dataBatch['normalPre'] normalPreBatch = Variable(normalPre_cpu).cuda() roughPre_cpu = dataBatch['roughPre'] roughPreBatch = Variable(roughPre_cpu).cuda() depthPre_cpu = dataBatch['depthPre'] depthPreBatch = Variable(depthPre_cpu).cuda() diffusePre_cpu = dataBatch['diffusePre'] diffusePreBatch = Variable(diffusePre_cpu).cuda() specularPre_cpu = dataBatch['specularPre'] specularPreBatch = Variable(specularPre_cpu).cuda() if albedoPreBatch.size(2) < opt.imHeight or albedoPreBatch.size( 3) < opt.imWidth: albedoPreBatch = F.interpolate(albedoPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if normalPreBatch.size(2) < opt.imHeight or normalPreBatch.size( 3) < opt.imWidth: normalPreBatch = F.interpolate(normalPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if roughPreBatch.size(2) < opt.imHeight or roughPreBatch.size( 3) < opt.imWidth: roughPreBatch = F.interpolate(roughPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if depthPreBatch.size(2) < opt.imHeight or depthPreBatch.size( 3) < opt.imWidth: depthPreBatch = F.interpolate(depthPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') # Regress the diffusePred and specular Pred envRow, envCol = diffusePreBatch.size(2), diffusePreBatch.size(3) imBatchSmall = F.adaptive_avg_pool2d(imBatch, (envRow, envCol)) diffusePreBatch, specularPreBatch = models.LSregressDiffSpec( diffusePreBatch.detach(), specularPreBatch.detach(), imBatchSmall, diffusePreBatch, specularPreBatch) if diffusePreBatch.size(2) < opt.imHeight or diffusePreBatch.size( 3) < opt.imWidth: diffusePreBatch = F.interpolate(diffusePreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if specularPreBatch.size(2) < opt.imHeight or specularPreBatch.size( 3) < opt.imWidth: specularPreBatch = F.interpolate(specularPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') envmapsPre_cpu = dataBatch['envmapsPre'] envmapsPreBatch = Variable(envmapsPre_cpu).cuda() # Normalize Albedo and depth bn, ch, nrow, ncol = albedoPreBatch.size() albedoPreBatch = albedoPreBatch.view(bn, -1) albedoPreBatch = albedoPreBatch / torch.clamp( torch.mean(albedoPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 albedoPreBatch = albedoPreBatch.view(bn, ch, nrow, ncol) bn, ch, nrow, ncol = depthPreBatch.size() depthPreBatch = depthPreBatch.view(bn, -1) depthPreBatch = depthPreBatch / torch.clamp( torch.mean(depthPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 depthPreBatch = depthPreBatch.view(bn, ch, nrow, ncol) ######################################################## # Build the cascade network architecture # if opt.cascadeLevel == 0: inputBatch = imBatch elif opt.cascadeLevel > 0: inputBatch = torch.cat([ imBatch, albedoPreBatch, normalPreBatch, roughPreBatch, depthPreBatch, diffusePreBatch, specularPreBatch ], dim=1) # Initial Prediction x1, x2, x3, x4, x5, x6 = encoder(inputBatch) albedoPred = 0.5 * (albedoDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) normalPred = normalDecoder(imBatch, x1, x2, x3, x4, x5, x6) roughPred = roughDecoder(imBatch, x1, x2, x3, x4, x5, x6) depthPred = 0.5 * (depthDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) albedoBatch = segBRDFBatch * albedoBatch albedoPred1 = models.LSregress( albedoPred.detach() * segBRDFBatch.expand_as(albedoPred), albedoBatch * segBRDFBatch.expand_as(albedoBatch), albedoPred) albedoPred1 = torch.clamp(albedoPred1, 0, 1) depthPred1 = models.LSregress( depthPred.detach() * segAllBatch.expand_as(depthPred), depthBatch * segAllBatch.expand_as(depthBatch), depthPred) ## Compute Errors pixelObjNum = (torch.sum(segBRDFBatch).cpu().data).item() pixelAllNum = (torch.sum(segAllBatch).cpu().data).item() albedoErr = torch.sum( (albedoPred1 - albedoBatch) * (albedoPred1 - albedoBatch) * segBRDFBatch.expand_as(albedoBatch) / pixelObjNum / 3.0) normalErr = torch.sum( (normalPred - normalBatch) * (normalPred - normalBatch) * segAllBatch.expand_as(normalBatch)) / pixelAllNum / 3.0 roughErr = torch.sum((roughPred - roughBatch) * (roughPred - roughBatch) * segBRDFBatch) / pixelObjNum depthErr = torch.sum( (torch.log(depthPred1 + 1) - torch.log(depthBatch + 1)) * (torch.log(depthPred1 + 1) - torch.log(depthBatch + 1)) * segAllBatch.expand_as(depthBatch)) / pixelAllNum return [albedoPred, albedoErr], [normalPred, normalErr], \ [roughPred, roughErr], [depthPred, depthErr]
def wrapperIIW(dataBatch, opt, encoder, albedoDecoder, normalDecoder, roughDecoder, depthDecoder): eq = dataBatch['eq'] darker = dataBatch['darker'] # Load the image from cpu to gpu im_cpu = dataBatch['im'] imBatch = Variable(im_cpu).cuda() if opt.cascadeLevel > 0: albedoPre_cpu = dataBatch['albedoPre'] albedoPreBatch = Variable(albedoPre_cpu).cuda() normalPre_cpu = dataBatch['normalPre'] normalPreBatch = Variable(normalPre_cpu).cuda() roughPre_cpu = dataBatch['roughPre'] roughPreBatch = Variable(roughPre_cpu).cuda() depthPre_cpu = dataBatch['depthPre'] depthPreBatch = Variable(depthPre_cpu).cuda() diffusePre_cpu = dataBatch['diffusePre'] diffusePreBatch = Variable(diffusePre_cpu).cuda() specularPre_cpu = dataBatch['specularPre'] specularPreBatch = Variable(specularPre_cpu).cuda() if albedoPreBatch.size(2) < opt.imHeight or albedoPreBatch.size( 3) < opt.imWidth: albedoPreBatch = F.interpolate(albedoPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if normalPreBatch.size(2) < opt.imHeight or normalPreBatch.size( 3) < opt.imWidth: normalPreBatch = F.interpolate(normalPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if roughPreBatch.size(2) < opt.imHeight or roughPreBatch.size( 3) < opt.imWidth: roughPreBatch = F.interpolate(roughPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if depthPreBatch.size(2) < opt.imHeight or depthPreBatch.size( 3) < opt.imWidth: depthPreBatch = F.interpolate(depthPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') # Regress the diffusePred and specular Pred envRow, envCol = diffusePreBatch.size(2), diffusePreBatch.size(3) imBatchSmall = F.adaptive_avg_pool2d(imBatch, (envRow, envCol)) diffusePreBatch, specularPreBatch = models.LSregressDiffSpec( diffusePreBatch.detach(), specularPreBatch.detach(), imBatchSmall, diffusePreBatch, specularPreBatch) if diffusePreBatch.size(2) < opt.imHeight or diffusePreBatch.size( 3) < opt.imWidth: diffusePreBatch = F.interpolate(diffusePreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') if specularPreBatch.size(2) < opt.imHeight or specularPreBatch.size( 3) < opt.imWidth: specularPreBatch = F.interpolate(specularPreBatch, [opt.imHeight, opt.imWidth], mode='bilinear') # Normalize Albedo and depth bn, ch, nrow, ncol = albedoPreBatch.size() albedoPreBatch = albedoPreBatch.view(bn, -1) albedoPreBatch = albedoPreBatch / torch.clamp( torch.mean(albedoPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 albedoPreBatch = albedoPreBatch.view(bn, ch, nrow, ncol) bn, ch, nrow, ncol = depthPreBatch.size() depthPreBatch = depthPreBatch.view(bn, -1) depthPreBatch = depthPreBatch / torch.clamp( torch.mean(depthPreBatch, dim=1), min=1e-10).unsqueeze(1) / 3.0 depthPreBatch = depthPreBatch.view(bn, ch, nrow, ncol) ######################################################## # Build the cascade network architecture # if opt.cascadeLevel == 0: inputBatch = imBatch elif opt.cascadeLevel > 0: inputBatch = torch.cat([ imBatch, albedoPreBatch, normalPreBatch, roughPreBatch, depthPreBatch, diffusePreBatch, specularPreBatch ], dim=1) # Initial Prediction x1, x2, x3, x4, x5, x6 = encoder(inputBatch) albedoPred = 0.5 * (albedoDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) normalPred = normalDecoder(imBatch, x1, x2, x3, x4, x5, x6) roughPred = roughDecoder(imBatch, x1, x2, x3, x4, x5, x6) depthPred = 0.5 * (depthDecoder(imBatch, x1, x2, x3, x4, x5, x6) + 1) eqLoss, darkerLoss = 0, 0 for m in range(0, albedoPred.size(0)): eqPoint = eq['point'][m, :].numpy().astype(np.long) eqWeight = eq['weight'][m, :].numpy().astype(np.float32) eqNum = eq['num'][m].numpy().astype(np.long) eqPoint = eqPoint[0:eqNum, :] eqWeight = eqWeight[0:eqNum] darkerPoint = darker['point'][m, :].numpy().astype(np.long) darkerWeight = darker['weight'][m, :].numpy().astype(np.float32) darkerNum = darker['num'][m].numpy().astype(np.long) darkerPoint = darkerPoint[0:darkerNum, :] darkerWeight = darkerWeight[0:darkerNum] eqL, darkerL = \ models.BatchRankingLoss(albedoPred[m, :], eqPoint, eqWeight, darkerPoint, darkerWeight ) eqLoss += eqL darkerLoss += darkerL eqLoss = eqLoss / max(albedoPred.size(0), 1e-5) darkerLoss = darkerLoss / max(albedoPred.size(0), 1e-5) return [albedoPred, None], [normalPred, None], \ [roughPred, None], [depthPred, None], \ [eq, eqLoss], [darker, darkerLoss]