all

Author: wwlCape

PLOT.py

@@ -0,0 +1,90 @@
+import matplotlib.pyplot as plt
+import pickle, glob
+import numpy as np
+import sys
+psnr_prefix = './psnr/*'
+psnr_paths = sorted(glob.glob(psnr_prefix))
+
+psnr_means = {}
+
+def filter_by_scale(row, scale):
+	return row[-1]==scale
+
+for i, psnr_path in enumerate(psnr_paths):
+	print ""
+	print psnr_path
+	psnr_dict = None
+	epoch = str(i)#psnr_path.split("_")[-1]
+	with open(psnr_path, 'rb') as f:
+		psnr_dict = pickle.load(f)
+	dataset_keys = psnr_dict.keys()
+	for j, key in enumerate(dataset_keys):
+		print 'dataset', key
+		psnr_list = psnr_dict[key]
+		psnr_np = np.array(psnr_list)
+
+		psnr_np_2 = psnr_np[np.array([filter_by_scale(row,2) for row in psnr_np])]
+		psnr_np_3 = psnr_np[np.array([filter_by_scale(row,3) for row in psnr_np])]
+		psnr_np_4 = psnr_np[np.array([filter_by_scale(row,4) for row in psnr_np])]
+		print "x2:",np.mean(psnr_np_2, axis=0).tolist()
+		print "x3:",np.mean(psnr_np_3, axis=0).tolist()
+		print "x4:",np.mean(psnr_np_4, axis=0).tolist()
+
+		mean_2 = np.mean(psnr_np_2, axis=0).tolist()
+		mean_3 = np.mean(psnr_np_3, axis=0).tolist()
+		mean_4 = np.mean(psnr_np_4, axis=0).tolist()
+		psnr_mean = [mean_2, mean_3, mean_4]
+		#print 'psnr mean', psnr_mean
+		if psnr_means.has_key(key):
+			psnr_means[key][epoch] = psnr_mean
+		else:
+			psnr_means[key] = {epoch: psnr_mean}
+
+#sys.exit(1)
+
+keys = psnr_means.keys()
+for i, key in enumerate(keys):
+	psnr_dict = psnr_means[key]
+	epochs = sorted(psnr_dict.keys())
+	x_axis = []
+	bicub_mean = []
+	vdsr_mean_2 = []
+	vdsr_mean_3 = []
+	vdsr_mean_4 = []
+
+	for epoch in epochs:
+		print epoch
+		print psnr_dict[epoch]
+		x_axis.append(int(epoch))
+		bicub_mean.append(psnr_dict[epoch][0][0])
+		vdsr_mean_2.append(psnr_dict[epoch][0][1])
+		vdsr_mean_3.append(psnr_dict[epoch][1][1])
+		vdsr_mean_4.append(psnr_dict[epoch][2][1])
+	plt.figure(i)
+	print key
+	print len(x_axis), len(bicub_mean), len(vdsr_mean_2)
+	print vdsr_mean_2
+	print "x2", np.argmax(vdsr_mean_2), np.max(vdsr_mean_2)
+	print "x3", np.argmax(vdsr_mean_3), np.max(vdsr_mean_3)
+	print "x4", np.argmax(vdsr_mean_4), np.max(vdsr_mean_4)
+	lines_bicub = plt.plot(vdsr_mean_2, 'g')
+	lines_bicub = plt.plot(vdsr_mean_4, 'b', vdsr_mean_3, 'y')
+	plt.setp(lines_bicub, linewidth=3.0)
+	plt.show()
+
+"""
+psnr_means :
+	{
+		'DATASET_NAME' :
+			{
+				'EPOCH' : [bicubic psnr, vdsr psnr]
+			}
+		'DATASET_NAME_2':
+			{
+				'EPOCH' : [bicubic psnr, vdsr psnr]
+			}
+		...
+	}
+"""
+# for i, psnr_path in enumerate(psnr_paths):
+	# print i, psnr_path

README.md

@@ -0,0 +1,2 @@
+# STAN
+This is an official implementation of Video Super-Resolution via a Spatio-Temporal Alignment Network

datasets.py

@@ -0,0 +1,112 @@
+import os, sys, math, random, glob, cv2
+import numpy as np
+### torch lib
+import torch
+import torch.utils.data as data
+### custom lib
+import utils
+import pdb
+import torchvision.transforms as transforms
+
+
+class RandomCrop(object):
+    def __init__(self, image_size, crop_size):
+        self.ch, self.cw = crop_size
+        ih, iw = image_size
+        self.h1 = random.randint(0, ih - self.ch)
+        self.w1 = random.randint(0, iw - self.cw)
+        self.h2 = self.h1 + self.ch
+        self.w2 = self.w1 + self.cw
+    def __call__(self, img):
+        if len(img.shape) == 3:
+            return img[self.h1 : self.h2, self.w1 : self.w2, :]
+        else:
+            return img[self.h1 : self.h2, self.w1 : self.w2]
+
+class MultiFramesDataset(data.Dataset):
+    def __init__(self, opts, mode):
+        super(MultiFramesDataset, self).__init__()
+        self.transform = transforms.Compose([transforms.ToTensor()])
+        self.opts = opts
+        self.mode = mode
+        self.task_videos = []
+        self.num_frames = []
+        self.dataset_task_list = []
+
+        list_filename = os.path.join(opts.list_dir, "train_tasks_%s.txt" %(opts.datasets_tasks))    
+        with open(list_filename) as f:
+            for line in f.readlines():
+                if line[0] != "#":
+                    self.dataset_task_list.append(line.strip().split())   
+        self.num_tasks = len(self.dataset_task_list)
+        for dataset, task in self.dataset_task_list:        
+            list_filename = os.path.join(opts.list_dir, "%s_%s.txt" %(dataset, mode))   
+            print("[%s] Read %s (Task %s)" %(self.__class__.__name__, list_filename, task))
+            with open(list_filename) as f:
+                videos = [line.rstrip() for line in f.readlines()] 
+            for video in videos:
+                self.task_videos.append([task, os.path.join(dataset, video)])
+                input_dir = os.path.join(self.opts.data_dir, self.mode, "input", dataset, video)    
+                frame_list = glob.glob(os.path.join(input_dir, '*.png'))
+                if len(frame_list) == 0:
+                    raise Exception("No frames in %s" %input_dir)
+                self.num_frames.append(len(frame_list))
+        print("[%s] Total %d videos (%d frames), %d tasks" %(self.__class__.__name__, len(self.task_videos), sum(self.num_frames), self.num_tasks))
+
+    def __len__(self):
+        return len(self.task_videos)
+    def __getitem__(self, index):
+        ## random select starting frame index t between [0, N - #sample_frames]
+        N = self.num_frames[index]
+        T = random.randint(0, N - self.opts.sample_frames)
+        task = self.task_videos[index][0]
+        video = self.task_videos[index][1]      
+        ## load input and processed frames
+        input_dir = os.path.join(self.opts.data_dir, self.mode, "input")
+        process_dir = os.path.join(self.opts.data_dir, self.mode, "processed", task)   
+        ## sample from T to T + #sample_frames - 1
+        frame_i = []
+        frame_p = []
+        frame_i_tmp = []
+        frame_p_tmp = []
+        for t in range(T, T + self.opts.sample_frames):
+            frame_i_tmp.append( utils.read_img(os.path.join(input_dir, video, "%08d.png" %t) ) )  
+            frame_p_tmp.append( utils.read_img(os.path.join(process_dir, video, "%08d.png" %t) ) )
+        ## data augmentation
+        if self.mode == 'train':
+            for t in range(self.opts.sample_frames):
+                frame_i_tmp[t], frame_p_tmp[t] = utils.get_patch(frame_i_tmp[t], frame_p_tmp[t], self.opts.crop_size, 4, multi_scale=False)
+                frame_i.append(frame_i_tmp[t])
+                frame_p.append(frame_p_tmp[t])
+            if self.opts.geometry_aug:
+                ### random rotate
+                rotate = random.randint(0, 3)
+                if rotate != 0:
+                    for t in range(self.opts.sample_frames):
+                        frame_i[t] = np.rot90(frame_i[t], rotate)
+                        frame_p[t] = np.rot90(frame_p[t], rotate)
+                ## horizontal flip
+                if np.random.random() >= 0.5:
+                    for t in range(self.opts.sample_frames):
+                        frame_i[t] = cv2.flip(frame_i[t], flipCode=0)
+                        frame_p[t] = cv2.flip(frame_p[t], flipCode=0)
+        elif self.mode == "test":
+            ## resize image to avoid size mismatch after downsampline and upsampling
+            H_i = frame_i[0].shape[0]
+            W_i = frame_i[0].shape[1]
+
+            H_o = int(math.ceil(float(H_i) / self.opts.size_multiplier) * self.opts.size_multiplier)
+            W_o = int(math.ceil(float(W_i) / self.opts.size_multiplier) * self.opts.size_multiplier)
+            for t in range(self.opts.sample_frames):
+                frame_i_tmp[t], frame_p_tmp[t] = utils.get_patch(frame_i_tmp[t], frame_p_tmp[t], self.opts.crop_size, 4, multi_scale=False)
+                frame_i.append(frame_i_tmp[t])
+                frame_p.append(frame_p_tmp[t])
+        
+        else:
+            raise Exception("Unknown mode (%s)" %self.mode)
+        ### convert (H, W, C) array to (C, H, W) tensor
+        data = []
+        for t in range(self.opts.sample_frames):
+            data.append(torch.from_numpy(frame_i[t].transpose(2, 0, 1).astype(np.float32)).contiguous())
+            data.append(torch.from_numpy(frame_p[t].transpose(2, 0, 1).astype(np.float32)).contiguous())
+        return data

install.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+cd ./networks/FAC/kernelconv2d
+python setup.py clean
+python setup.py install --user

lists/REDS4_test.txt

@@ -0,0 +1,4 @@
+000
+011
+015
+020

lists/Vid4_test.txt

@@ -0,0 +1,4 @@
+calendar
+city
+foliage
+walk