search

ryujaehun / torch_sr

torch_sr
1 stars 0 forks source link

Video #5

Open ryujaehun opened 6 years ago

ryujaehun commented 6 years ago

https://github.com/ryujaehun/ESPCN https://github.com/ryujaehun/video-super-resolution

ryujaehun commented 6 years ago

https://people.csail.mit.edu/celiu/CVPR2011/default.html

ryujaehun commented 6 years ago

https://pixabay.com/en/videos/list/

ryujaehun commented 6 years ago

change codec

import argparse
import os
from os import listdir

import cv2
import numpy as np
import torch
from PIL import Image
from torch.autograd import Variable
from torchvision.transforms import ToTensor
from tqdm import tqdm

from utils.dataset import is_video_file

#from model import Net

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test Super Resolution')
    parser.add_argument('--upscale_factor', default=2, type=int, help='super resolution upscale factor')
    parser.add_argument('--is_real_time','-r', default=False, type=bool, help='super resolution real time to show')
    parser.add_argument('--delay_time', default=1, type=int, help='super resolution delay time to show')
    parser.add_argument('--model_name','-m', default='model_epoch_80.pth', type=str, help='super resolution model name')
    opt = parser.parse_args()

    UPSCALE_FACTOR = opt.upscale_factor
    IS_REAL_TIME = opt.is_real_time
    DELAY_TIME = opt.delay_time
    MODEL_NAME = opt.model_name

    path = 'dataset/video/'
    videos_name = [x for x in listdir(path) if is_video_file(x)]

    model=torch.load(opt.model_name)
    model=model.cuda()
    #model = Net(upscale_factor=UPSCALE_FACTOR)
    #if torch.cuda.is_available():
    #        model = model.cuda()
    # for cpu
    # model.load_state_dict(torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage))
    #model.load_state_dict(torch.load('epochs/' + MODEL_NAME))

    out_path = 'result/video/'+str(opt.upscale_factor)+'/'
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    for video_name in tqdm(videos_name, desc='convert LR videos to HR videos'):
        print(path + video_name)
        videoCapture = cv2.VideoCapture(path + video_name)
        if not IS_REAL_TIME:
            fps = videoCapture.get(cv2.CAP_PROP_FPS)
            size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH) * UPSCALE_FACTOR),
                    int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)) * UPSCALE_FACTOR)
            output_name = out_path + video_name.split('.')[0] + '.mp4'
            videoWriter = cv2.VideoWriter(output_name,int(videoCapture.get(cv2.CAP_PROP_FOURCC)), fps, size)
        # read frame
        success, frame = videoCapture.read()
        while success:
            img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).convert('YCbCr')
            y, cb, cr = img.split()
            image = Variable(ToTensor()(y)).view(1, -1, y.size[1], y.size[0])
            if torch.cuda.is_available():
                image = image.cuda()

            out = model(image)
            out = out.cpu()
            out_img_y = out.data[0].numpy()
            out_img_y *= 255.0
            out_img_y = out_img_y.clip(0, 255)
            out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')
            out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)
            out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
            out_img = Image.merge('YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')
            out_img = cv2.cvtColor(np.asarray(out_img), cv2.COLOR_RGB2BGR)

            if IS_REAL_TIME:
                cv2.imshow('LR Video ', frame)
                cv2.imshow('SR Video ', out_img)
                cv2.waitKey(DELAY_TIME)
            else:
                # save video
                videoWriter.write(out_img)
            # next frame
            success, frame = videoCapture.read()
ryujaehun commented 6 years ago

cv logic

import argparse
import os
from os import listdir

import cv2
import numpy as np
import torch
from PIL import Image
from torch.autograd import Variable
from torchvision.transforms import ToTensor
from tqdm import tqdm

from utils.dataset import is_video_file

#from model import Net

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test Super Resolution')
    parser.add_argument('--upscale_factor', default=2, type=int, help='super resolution upscale factor')
    parser.add_argument('--is_real_time','-r', default=False, type=bool, help='super resolution real time to show')
    parser.add_argument('--delay_time', default=1, type=int, help='super resolution delay time to show')
    parser.add_argument('--model_name','-m', default='model_epoch_80.pth', type=str, help='super resolution model name')
    opt = parser.parse_args()

    UPSCALE_FACTOR = opt.upscale_factor
    IS_REAL_TIME = opt.is_real_time
    DELAY_TIME = opt.delay_time
    MODEL_NAME = opt.model_name

    path = 'dataset/video/'
    videos_name = [x for x in listdir(path) if is_video_file(x)]

    model=torch.load(opt.model_name)
    model=model.cuda()
    #model = Net(upscale_factor=UPSCALE_FACTOR)
    #if torch.cuda.is_available():
    #        model = model.cuda()
    # for cpu
    # model.load_state_dict(torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage))
    #model.load_state_dict(torch.load('epochs/' + MODEL_NAME))

    out_path = 'result/video/'+str(opt.upscale_factor)+'/'
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    for video_name in tqdm(videos_name, desc='convert LR videos to HR videos'):
        print(path + video_name)
        videoCapture = cv2.VideoCapture(path + video_name)
        if not IS_REAL_TIME:
            fps = videoCapture.get(cv2.CAP_PROP_FPS)
            size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH) * UPSCALE_FACTOR),
                    int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)) * UPSCALE_FACTOR)
            output_name = out_path + video_name.split('.')[0] + '.mp4'
            videoWriter = cv2.VideoWriter(output_name,int(videoCapture.get(cv2.CAP_PROP_FOURCC)), fps, size)
        # read frame
        success, frame = videoCapture.read()
        frame=cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
        y=frame[:,:,0]
        y=y.astype(float)
        y=y/255.
        while success:
            image = Variable(torch.from_numpy(y).float()).view(1, -1, frame.shape[0], frame.shape[1])
            if torch.cuda.is_available():
                image = image.cuda()

            out = model(image)
            out = out.cpu()
            out_img_y = out.data[0].numpy().astype(np.float32)

            out_img_y = out_img_y*255.
            out_img_y[out_img_y<0] = 0
            out_img_y[out_img_y>255.] = 255.
            out_img_y = out_img_y[0,:,:]
            out_img_y *= 255.0
            out_img_y=out_img_y.astype(np.uint8)
            out_cb=cv2.resize(frame[:,:,1],(frame.shape[1]*2,frame.shape[0]*2),interpolation=cv2.INTER_CUBIC)
            out_cr=cv2.resize(frame[:,:,2],(frame.shape[1]*2,frame.shape[0]*2),interpolation=cv2.INTER_CUBIC)
            out_img=np.stack((out_img_y,out_cb,out_cr),axis=2)

            out_img=cv2.cvtColor(out_img,cv2.COLOR_YCrCb2BGR)

            if IS_REAL_TIME:
                cv2.imshow('LR Video ', frame)
                cv2.imshow('SR Video ', out_img)
                cv2.waitKey(DELAY_TIME)
            else:
                # save video
                videoWriter.write(out_img)
            # next frame
            success, frame = videoCapture.read()
ryujaehun commented 6 years ago 
#! /usr/bin/python3
import cv2
import os
import glob
root=os.getcwd()
folder='video'
vidiofold=os.path.join(root,folder)
videolist=glob.glob(vidiofold+'/*.mp4')
for i in videolist:
    vidcap = cv2.VideoCapture(i)
    name=i.split('/')[-1].split('-')[0].strip()
    os.makedirs(os.path.join(root,name))
    with open(os.path.join(os.path.join(root,name),'meta.txt'), 'w') as f:
        f.write(str(int(vidcap.get(cv2.CAP_PROP_FOURCC)))+'\n')
        f.write(str(vidcap.get(cv2.CAP_PROP_FPS))+'\n')
    success,image = vidcap.read()

    count = 0
    success = True
    while success:
        cv2.imwrite(os.path.join(os.path.join(root,name),"%05d.bmp" % count), image)     # save frame as PNG file
        success,image = vidcap.read()
        #print('Read a new frame: ', success)
        count += 1
ryujaehun commented 6 years ago

merge

#! /usr/bin/python3
import cv2
import os,glob
root=os.getcwd()
join=os.path.join
dataset=join(root,'dataset')
folder=join(root,'vresult')
image_list=[img for img in os.listdir(folder)]
for i in image_list:
    temp=join(dataset,i)
    temp=join(temp,'meta.txt')
    with open(temp,'r') as f:
        foc=f.readline()
        fps=f.readline()
    #print(fps.strip(),foc.strip())
    im_list=glob.glob(join(folder,i+"/*bmp"))
    #print(type(im_list))
    im_list.sort()
    #print(im_list)
    frame = cv2.imread(im_list[0])
    height, width, layers = frame.shape
    fourcc =foc.strip()
    video = cv2.VideoWriter(join(folder,i)+'.mp4',int(fourcc),float(fps.strip()), (width,height))
    for image in im_list:
        video.write(cv2.imread(image))
    cv2.destroyAllWindows()
    video.release()
ryujaehun commented 6 years ago

srVideo.py

from __future__ import print_function
import argparse
import torch
from torch.autograd import Variable
from PIL import Image
from torchvision.transforms import ToTensor
import os,glob
import numpy as np
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
parser.add_argument('--input','-i',default='Animals' ,type=str, required=False, help='input image to use')
parser.add_argument('--model', type=str, default='model_epoch_80.pth',required=False, help='model file to use')
parser.add_argument('--cuda', action='store_true', required=False, help='use cuda')
opt = parser.parse_args()
join=os.path.join
root=os.getcwd()
result=join(root,'vresult')
result=join(result,opt.input)
dataset=join(root,'dataset')
datasetfolder=join(dataset,opt.input)
videolist=glob.glob(datasetfolder+'/*.bmp')
videolist.sort()
print(opt)
if not os.path.exists(result):
    os.makedirs(result)
for i in videolist:
    img = Image.open(join(datasetfolder,i)).convert('YCbCr')
    y, cb, cr = img.split()
    model = torch.load(opt.model)
    input = Variable(ToTensor()(y)).view(1, -1, y.size[1], y.size[0])
    if opt.cuda:
        model = model.cuda()
        input = input.cuda()
    out = model(input)
    out = out.cpu()
    out_img_y = out.data[0].numpy()
    out_img_y *= 255.0
    out_img_y = out_img_y.clip(0, 255)
    out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')
    out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)
    out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
    out_img = Image.merge('YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')
    out_img.save(join(result,i.split('/')[-1]))
ryujaehun commented 6 years ago

SR 에서 노이즈가 생기는 것을 모든 frame을 bmp로 쪼개서 SR함 그것을 merge

ryujaehun commented 6 years ago 
#! /usr/bin/python3
import cv2
import os
import glob
root=os.getcwd()
folder='video'
vidiofold=os.path.join(root,folder)
videolist=glob.glob(vidiofold+'/*.mp4')
for i in videolist:
    vidcap = cv2.VideoCapture(i)
    name=i.split('/')[-1].split('-')[0].strip()
    os.makedirs(os.path.join(root,name))
    with open(os.path.join(os.path.join(root,name),'meta.txt'), 'w') as f:
        f.write(str(int(vidcap.get(cv2.CAP_PROP_FOURCC)))+'\n')
        f.write(str(vidcap.get(cv2.CAP_PROP_FPS))+'\n')
    success,image = vidcap.read()

    count = 0
    success = True
    while success:
        cv2.imwrite(os.path.join(os.path.join(root,name),"%05d.bmp" % count), image)     # save frame as PNG file
        success,image = vidcap.read()
        #print('Read a new frame: ', success)
        count += 1