add heatmap code

SimpleNet/test.py

@@ -1,5 +1,19 @@
+import cv2
+import numpy as np
+from utils import *
+from tqdm import tqdm
+from loss import *
+from dataloader import TestLoader, SaliconDataset
+from scipy.stats import multivariate_normal
+import torch.nn.functional as F
+import torch.nn.init as init
+from torch.utils.data import DataLoader
+from PIL import Image
+from torchvision import transforms, utils
+import matplotlib.pyplot as plt
 import argparse
-import glob, os
+import glob
+import os
 import torch
 import sys
 import time
@@ -9,31 +23,20 @@
 from torch.autograd import Variable
 import matplotlib
 matplotlib.use('Agg')
-import matplotlib.pyplot as plt
-from torchvision import transforms, utils
-from PIL import Image
-from torch.utils.data import DataLoader
-import numpy as np, cv2
-import torch
-import torch.nn as nn
-import torch.nn.init as init
-import torch.nn.functional as F
-from scipy.stats import multivariate_normal
-from dataloader import TestLoader, SaliconDataset
-from loss import *
-from tqdm import tqdm
-from utils import *
 
 parser = argparse.ArgumentParser()
 
-parser.add_argument('--val_img_dir',default="../images/", type=str)
-parser.add_argument('--model_val_path',default="../saved_models/salicon_pnas.pt", type=str)
-parser.add_argument('--no_workers',default=4, type=int)
-parser.add_argument('--enc_model',default="pnas", type=str)
-parser.add_argument('--results_dir',default="../results/", type=str)
-parser.add_argument('--validate',default=0, type=int)
-parser.add_argument('--save_results',default=1, type=int)
-parser.add_argument('--dataset_dir',default="/home/samyak/old_saliency/saliency/SALICON_NEW/", type=str)
+parser.add_argument('--val_img_dir', default="../images/", type=str)
+parser.add_argument('--model_val_path',
+                    default="../saved_models/salicon_pnas.pt", type=str)
+parser.add_argument('--no_workers', default=4, type=int)
+parser.add_argument('--enc_model', default="pnas", type=str)
+parser.add_argument('--results_dir', default="../results/", type=str)
+parser.add_argument('--heatmap_dir', default="../heatmap_results/", type=str)
+parser.add_argument('--validate', default=0, type=int)
+parser.add_argument('--save_results', default=1, type=int)
+parser.add_argument(
+    '--dataset_dir', default="/home/samyak/old_saliency/saliency/SALICON_NEW/", type=str)
 
 args = parser.parse_args()
 
@@ -54,7 +57,7 @@
     print("ResNet Model")
     from model import ResNetModel
     model = ResNetModel()
-    
+
 elif args.enc_model == "vgg":
     print("VGG Model")
     from model import VGGModel
@@ -65,15 +68,17 @@
     from model import MobileNetV2
     model = MobileNetV2()
 
-if args.enc_model!="mobilenet":
+if args.enc_model != "mobilenet":
     model = nn.DataParallel(model)
 model.load_state_dict(torch.load(args.model_val_path))
 
 model = model.to(device)
 
 val_img_ids = os.listdir(args.val_img_dir)
 val_dataset = TestLoader(args.val_img_dir, val_img_ids)
-vis_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=args.no_workers)
+vis_loader = torch.utils.data.DataLoader(
+    val_dataset, batch_size=1, shuffle=False, num_workers=args.no_workers)
+
 
 def validate(model, loader, device, args):
     model.eval()
@@ -83,37 +88,41 @@ def validate(model, loader, device, args):
     kldiv_loss = AverageMeter()
     nss_loss = AverageMeter()
     sim_loss = AverageMeter()
-    
+
     for (img, gt, fixations) in tqdm(loader):
         img = img.to(device)
         gt = gt.to(device)
         fixations = fixations.to(device)
-        
+
         pred_map = model(img)
 
         # Blurring
         blur_map = pred_map.cpu().squeeze(0).clone().numpy()
         blur_map = blur(blur_map).unsqueeze(0).to(device)
-
-        cc_loss.update(cc(blur_map, gt))    
-        kldiv_loss.update(kldiv(blur_map, gt))    
-        nss_loss.update(nss(blur_map, gt))    
-        sim_loss.update(similarity(blur_map, gt))    
 
-    print('CC : {:.5f}, KLDIV : {:.5f}, NSS : {:.5f}, SIM : {:.5f}  time:{:3f} minutes'.format(cc_loss.avg, kldiv_loss.avg, nss_loss.avg, sim_loss.avg, (time.time()-tic)/60))
+        cc_loss.update(cc(blur_map, gt))
+        kldiv_loss.update(kldiv(blur_map, gt))
+        nss_loss.update(nss(blur_map, gt))
+        sim_loss.update(similarity(blur_map, gt))
+
+    print('CC : {:.5f}, KLDIV : {:.5f}, NSS : {:.5f}, SIM : {:.5f}  time:{:3f} minutes'.format(
+        cc_loss.avg, kldiv_loss.avg, nss_loss.avg, sim_loss.avg, (time.time()-tic)/60))
     sys.stdout.flush()
-    
+
     return cc_loss.avg
 
+
 if args.validate:
-	val_img_dir = args.dataset_dir + "images/val/"
-	val_gt_dir = args.dataset_dir + "maps/val/"
-	val_fix_dir = args.dataset_dir + "fixations/fixations/"
-
-	val_img_ids = [nm.split(".")[0] for nm in os.listdir(val_img_dir)]
-	val_dataset = SaliconDataset(val_img_dir, val_gt_dir, val_fix_dir, val_img_ids)
-	val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=args.no_workers)
-	with torch.no_grad():
-		validate(model, val_loader, device, args)
+    val_img_dir = args.dataset_dir + "images/val/"
+    val_gt_dir = args.dataset_dir + "maps/val/"
+    val_fix_dir = args.dataset_dir + "fixations/fixations/"
+
+    val_img_ids = [nm.split(".")[0] for nm in os.listdir(val_img_dir)]
+    val_dataset = SaliconDataset(
+        val_img_dir, val_gt_dir, val_fix_dir, val_img_ids)
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset, batch_size=1, shuffle=False, num_workers=args.no_workers)
+    with torch.no_grad():
+        validate(model, val_loader, device, args)
 if args.save_results:
-	visualize_model(model, vis_loader, device, args)
+    visualize_model(model, vis_loader, device, args)

SimpleNet/utils.py

@@ -1,4 +1,5 @@
-import cv2, os
+import cv2
+import os
 import torch
 from os.path import join
 import matplotlib.pyplot as plt
@@ -7,16 +8,18 @@
 from tqdm import tqdm
 from PIL import Image
 
+
 def blur(img):
     k_size = 11
-    bl = cv2.GaussianBlur(img,(k_size,k_size),0)
+    bl = cv2.GaussianBlur(img, (k_size, k_size), 0)
     return torch.FloatTensor(bl)
 
+
 def plot(pred, gt, args, idx):
     pred_npimg = utils.make_grid(pred.cpu()).numpy()
     gt_npimg = utils.make_grid(gt.cpu()).numpy()
 
-    plt.figure(figsize=(10,10))
+    plt.figure(figsize=(10, 10))
     plt.subplot(121)
     plt.title("Original")
     plt.imshow(np.transpose(gt_npimg, (1, 2, 0)))
@@ -25,38 +28,55 @@ def plot(pred, gt, args, idx):
     plt.title("Predicted")
     plt.imshow(np.transpose(pred_npimg, (1, 2, 0)))
 
-
     plt.savefig(args.results_dir + '{}_{}.png'.format(epoch, idx+1))
 
+
 def visualize_model(model, loader, device, args):
     with torch.no_grad():
         model.eval()
         os.makedirs(args.results_dir, exist_ok=True)
-        
+
         for (img, img_id, sz) in tqdm(loader):
             img = img.to(device)
-            
+
             pred_map = model(img)
             pred_map = pred_map.cpu().squeeze(0).numpy()
             pred_map = cv2.resize(pred_map, (sz[0], sz[1]))
-            
+
             pred_map = torch.FloatTensor(blur(pred_map))
-            img_save(pred_map, join(args.results_dir, img_id[0]), normalize=True)
+            img_save(pred_map, join(
+                args.results_dir, img_id[0]), normalize=True)
+
+            heat_map = im2heat(join(args.results_dir, img_id[0]), join(
+                args.val_img_dir, img_id[0]))
+
+            im = Image.fromarray(heat_map)
+            exten = join(args.heatmap_dir, img_id[0]).split('.')[-1]
+            if exten == "png":
+                im.save(join(args.heatmap_dir,
+                             img_id[0]), format=None, compress_level=0)
+            else:
+                # for jpg
+                im.save(join(args.heatmap_dir,
+                             img_id[0]), format=None, quality=100)
+
 
 def img_save(tensor, fp, nrow=8, padding=2,
-               normalize=False, range=None, scale_each=False, pad_value=0, format=None):
+             normalize=False, range=None, scale_each=False, pad_value=0, format=None):
     grid = utils.make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
-                     normalize=normalize, range=range, scale_each=scale_each)
+                           normalize=normalize, range=range, scale_each=scale_each)
 
     ''' Add 0.5 after unnormalizing to [0, 255] to round to nearest integer '''
-
-    ndarr = torch.round(grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0)).to('cpu', torch.uint8).numpy()
+
+    ndarr = torch.round(grid.mul(255).add_(0.5).clamp_(
+        0, 255).permute(1, 2, 0)).to('cpu', torch.uint8).numpy()
     im = Image.fromarray(ndarr)
     exten = fp.split('.')[-1]
-    if exten=="png":
+    if exten == "png":
         im.save(fp, format=format, compress_level=0)
     else:
-        im.save(fp, format=format, quality=100) #for jpg
+        im.save(fp, format=format, quality=100)  # for jpg
+
 
 class AverageMeter(object):
 
@@ -72,18 +92,23 @@ def reset(self):
         self.sum = 0
         self.count = 0
 
-    def update(self, val, n = 1):
+    def update(self, val, n=1):
         self.val = val
         self.sum += val*n
         self.count += n
         self.avg = self.sum / self.count
 
-def im2heat(pred_dir, a, gt, exten='.png'):
-    pred_nm = pred_dir + a + exten
-    pred = cv2.imread(pred_nm, 0)
+
+def convert(img):
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+
+def im2heat(map_path, rgb_image_path):
+    gt = cv2.imread(rgb_image_path)
+    gt = convert(gt)
+    pred = cv2.imread(map_path, 0)
     heatmap_img = cv2.applyColorMap(pred, cv2.COLORMAP_JET)
     heatmap_img = convert(heatmap_img)
-    pred = np.stack((pred, pred, pred),2).astype('float32')
+    pred = np.stack((pred, pred, pred), 2).astype('float32')
     pred = pred / 255.0
-
-    return np.uint8(pred * heatmap_img + (1.0-pred) * gt)
+    return np.uint8(pred * heatmap_img + (1.0-pred) * gt)