fixed image tensorboard viewer

Author: charlio23

datasets.py

@@ -39,10 +39,10 @@ def __getitem__(self, i):
 
         inputImage = rgb2lab(np.array(self.transf(img))).astype(float)
         inputImage[:,:,0] -= 50.0
-        image_ab = self.toTensor(cv2.resize(inputImage, (56, 56), interpolation = cv2.INTER_AREA)[:,:,1:].astype(float))
-        image_L = torch.from_numpy(inputImage[:,:,0]).unsqueeze_(0)
+        image_ab = self.toTensor(cv2.resize(inputImage, (56, 56), interpolation = cv2.INTER_AREA)[:,:,1:].astype(float)).float()
+        image_L = torch.from_numpy(inputImage[:,:,0]).unsqueeze_(0).float()
         if self.output_full:
-            img_all = torch.from_numpy(inputImage)
+            img_all = self.toTensor(inputImage).float()
             return img_all, image_L, image_ab
         else:
             return image_L, image_ab

train_coGAN.py

@@ -143,7 +143,7 @@ def train(train_loader, model_G, model_D, criterion_G, criterion_GAN, optimizer_
         data_time.update(time.time() - end)
         var = Variable(img_L.float(), requires_grad=True).cuda()
         real = Variable(real.float(), requires_grad=True).cuda()
-        target = Variable(utils.soft_encode_ab(target).float(), requires_grad=False).cuda()
+        target_class = Variable(utils.soft_encode_ab(target).float(), requires_grad=False).cuda()
         # compute output G(L)
         output = model_G(var)
 
@@ -174,7 +174,7 @@ def train(train_loader, model_G, model_D, criterion_G, criterion_GAN, optimizer_
         # Fool the discriminator
         loss_G_GAN = criterion_GAN(fake_prob, True)
         # Regressor loss term
-        loss_G_L2 = criterion_G(output, target)
+        loss_G_L2 = criterion_G(output, target_class)
         loss_G = loss_G_GAN + loss_G_L2*10
         if torch.isnan(loss_G):
             print('NaN value encountered in loss_G.')
@@ -208,7 +208,7 @@ def train(train_loader, model_G, model_D, criterion_G, criterion_GAN, optimizer_
             start = time.time()
             batch_num = np.maximum(args.batch_size//4,2)
             idx = i + epoch*len(train_loader)
-            imgs = utils.getImages(img_L.float(), target.cpu(), output.detach().cpu(), batch_num, decode=True)
+            imgs = utils.getImages(var.detach(), target.cuda(), output.detach(), batch_num, do_decode=True)
             writer.add_image('data/imgs_gen', imgs, idx)
             print("Img conversion time: ", time.time() - start)
         writer.add_scalar('data/L2_loss_train', losses_L2.avg, i + epoch*len(train_loader))

utils.py

@@ -62,18 +62,18 @@ def generateImg(Z,light):
     Image = cv2.cvtColor(newImg.astype(np.uint8), cv2.COLOR_LAB2RGB)
     return Image
 
-def getImages(L_channel, ab_target, ab_gen, batch_num, decode=True):
+def getImages(L_channel, ab_target, ab_gen, batch_num, do_decode=True):
     L_channel = interpolate(L_channel[:batch_num,:,:,:] + 50.0, scale_factor=0.25, mode='bilinear', 
         recompute_scale_factor=True, align_corners=True)
-    if decode:
+    if do_decode:
         ab_gen = decode(ab_gen[:batch_num,:,:,:], T=0.38)
     else:
         ab_gen = ab_gen[:batch_num,:,:,:]
 
     ab_target = ab_target[:batch_num,:,:,:]
     img_target = torch.cat([L_channel, ab_target], dim=1)
     img_gen = torch.cat([L_channel, ab_gen], dim=1)
-    img_all = torch.cat([img_target, img_gen], dim=0).numpy().transpose((0,2,3,1))
+    img_all = torch.cat([img_target, img_gen], dim=0).cpu().numpy().transpose((0,2,3,1))
 
     imgs_all_l = []
     for i in range(batch_num):