EDVR
EDVR copied to clipboard
Published
20 hours ago •
xinntao
xinntao
What can be reason why I get this error: TypeError: DCNv2Pack() takes no arguments ?
I'm trying to use this function:
import ffmpeg
import os
import os.path as osp
import glob
import logging
import cv2
import torch
import numpy as np
import shutil
import re
import gc
import youtube_dl
from PIL import Image
from pathlib import Path
from tqdm import tqdm
import basicsr.utils as util
import basicsr.data.data_util as data_util
import basicsr.models.archs.edvr_arch as EDVR_arch
workfolder = Path('./video')
source_folder = workfolder / "source"
inframes_root = workfolder / "inframes"
audio_root = workfolder / "audio"
outframes_root = workfolder / "outframes"
result_folder = workfolder / "result"
pretrained_models = Path('../experiments/pretrained_models')
def clean_mem():
# torch.cuda.empty_cache()
gc.collect()
def get_fps(source_path: Path) -> str:
print(source_path)
probe = ffmpeg.probe(str(source_path))
stream_data = next(
(stream for stream in probe['streams'] if stream['codec_type'] == 'video'),
None,
)
return stream_data['avg_frame_rate']
def download_video_from_url(source_url, source_path: Path, quality: str):
if source_path.exists():
source_path.unlink()
ydl_opts = {
'format': 'bestvideo[height<={}][ext=mp4]+bestaudio[ext=m4a]/mp4'.format(quality),
'outtmpl': str(source_path),
}
with youtube_dl.YoutubeDL(ydl_opts) as ydl:
ydl.download([source_url])
def preProcess(imag_path_l, multiple):
'''Need to resize images for blurred model (needs to be multiples of 16)'''
for img_path in imag_path_l:
im = Image.open(img_path)
h, w = im.size
# resize so they are multiples of 4 or 16 (for blurred)
h = h - h % multiple
w = w - w % multiple
im = im.resize((h,w))
im.save(img_path)
def purge_images(dir):
for f in os.listdir(dir):
if re.search('.*?\.jpg', f):
os.remove(os.path.join(dir, f))
def extract_raw_frames(source_path: Path):
inframes_folder = inframes_root / (source_path.stem)
inframe_path_template = str(inframes_folder / '%5d.jpg')
inframes_folder.mkdir(parents=True, exist_ok=True)
purge_images(inframes_folder)
ffmpeg.input(str(source_path)).output(
str(inframe_path_template), format='image2', vcodec='mjpeg', qscale=0
).run(capture_stdout=True)
def make_subfolders(img_path_l, chunk_size):
i = 0
subFolderList = []
source_img_path = Path('/content/EDVR/codes/video/inframes/video_subfolders')
source_img_path.mkdir(parents=True, exist_ok=True)
for img in img_path_l:
if i % chunk_size == 0:
img_path = source_img_path / str(i)
img_path.mkdir(parents=True, exist_ok=True)
subFolderList.append(str(img_path))
i+=1
img_name = osp.basename(img)
img_path_name = img_path / img_name
shutil.copyfile(img, img_path_name)
return subFolderList
def remove_subfolders():
shutil.rmtree('/content/EDVR/codes/video/inframes/video_subfolders', ignore_errors=True, onerror=None)
def edvrPredict(data_mode, chunk_size, stage):
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_mode = data_mode # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = stage # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
model_path = pretrained_models / 'EDVR_Vimeo90K_SR_L.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
elif data_mode == 'sharp_bicubic':
if stage == 1:
model_path = pretrained_models / 'EDVR_REDS_SR_L.pth'
else:
model_path = pretrained_models / 'EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = pretrained_models / 'EDVR_REDS_SRblur_L.pth'
else:
model_path = pretrained_models / 'EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = pretrained_models / 'EDVR_REDS_deblur_L.pth'
else:
model_path = pretrained_models / 'EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = pretrained_models / 'EDVR_REDS_deblurcomp_L.pth'
else:
model_path = pretrained_models / 'EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
print('Model Used: ', model_path)
if data_mode == 'Vid4':
N_in = 7 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 40
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
if data_mode == 'TOF':
model = TOF_arch.TOFlow(adapt_official=True)
else:
model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, with_predeblur=predeblur, hr_in=HR_in)
#### dataset
test_dataset_folder = '/content/EDVR/codes/video/inframes'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode in ('Vid4','sharp_bicubic'):
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '/content/EDVR/codes/video/outframes'
util.mkdirs(save_folder)
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
# remove old video_subfolder if exists
remove_subfolders()
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
# for each subfolder
for subfolder in subfolder_l:
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
if save_imgs:
util.mkdirs(save_subfolder)
purge_images(save_subfolder)
# preprocess images (needed for blurred models)
if predeblur:
preProcess(img_path_l, 16)
else:
preProcess(img_path_l, 4)
# make even more subfolders
subFolderList = make_subfolders(img_path_l, chunk_size)
#### read LQ and GT images in chunks of 1000
for subSubFolder in subFolderList:
clean_mem()
imgs_LQ = data_util.read_img_seq(subSubFolder)
subSubFolder_l = sorted(glob.glob(osp.join(subSubFolder, '*')))
max_idx = len(subSubFolder_l)
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in tqdm(enumerate(subSubFolder_l)):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(osp.join(save_subfolder, '{}.jpg'.format(img_name)), output)
# print('Saved Image:', str(osp.join(save_subfolder, '{}.jpg'.format(img_name))))
def moveProcessedFrames():
shutil.rmtree('/content/EDVR/codes/video/inframes')
os.rename('/content/EDVR/codes/video/outframes', '/content/EDVR/codes/video/inframes')
def build_video(source_path: Path) -> Path:
out_path = result_folder / (
source_path.name.replace('.mp4', '_no_audio.mp4')
)
outframes_folder = outframes_root / (source_path.stem)
outframes_path_template = str(outframes_folder / '%5d.jpg')
out_path.parent.mkdir(parents=True, exist_ok=True)
if out_path.exists():
out_path.unlink()
fps = get_fps(source_path)
print('Original FPS is: ', fps)
ffmpeg.input(
str(outframes_path_template),
format='image2',
vcodec='mjpeg',
framerate=fps,
).output(str(out_path), crf=17, vcodec='libx264').run(capture_stdout=True)
result_path = result_folder / source_path.name
if result_path.exists():
result_path.unlink()
# making copy of non-audio version in case adding back audio doesn't apply or fails.
shutil.copyfile(str(out_path), str(result_path))
# adding back sound here
audio_file = Path(str(source_path).replace('.mp4', '.aac'))
if audio_file.exists():
audio_file.unlink()
os.system(
'ffmpeg -y -i "'
+ str(source_path)
+ '" -vn -acodec copy "'
+ str(audio_file)
+ '"'
)
if audio_file.exists:
os.system(
'ffmpeg -y -i "'
+ str(out_path)
+ '" -i "'
+ str(audio_file)
+ '" -shortest -c:v copy -c:a aac -b:a 256k "'
+ str(result_path)
+ '"'
)
print('Video created here: ' + str(result_path))
return result_path
def edvr_video(source_url: str, source_path: Path, data_mode: str, original_quality: str,
chunk_size: int, finetune_stage2: bool):
# download video
download_video_from_url(source_url , source_path, original_quality)
# extract frames
extract_raw_frames(source_path)
# process frames
edvrPredict(data_mode, chunk_size, 1)
# fine-tune stage 2
if finetune_stage2:
# move the stage 1 processed frames over
moveProcessedFrames()
# process again
edvrPredict(data_mode, chunk_size, 2)
# build back video
build_video(source_path)
But when I try to implement that function I get this error and I dint understand why TypeError: DCNv2Pack() takes no arguments. What can be reason? Thanks
model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, with_predeblur=predeblur, hr_in=HR_in)
Next issue is here :
[youtube] v6i3uccnZhQ: Downloading webpage [download] Destination: /content/video.f133.mp4 [download] 100% of 1.75MiB in 00:35 [download] Destination: /content/video.mp4.f140 [download] 100% of 943.22KiB in 00:18 [ffmpeg] Merging formats into "/content/video.mp4" Deleting original file /content/video.f133.mp4 (pass -k to keep) Deleting original file /content/video.mp4.f140 (pass -k to keep) Model Used: ../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth
TypeError Traceback (most recent call last)
3 frames
/content/EDVR/basicsr/models/archs/edvr_arch.py in init(self, num_feat, deformable_groups) 50 3, 51 padding=1, ---> 52 deformable_groups=deformable_groups) 53 54 if i < 3:
TypeError: DCNv2Pack() takes no arguments
