xinntao / EDVR

Winning Solution in NTIRE19 Challenges on Video Restoration and Enhancement (CVPR19 Workshops) - Video Restoration with Enhanced Deformable Convolutional Networks. EDVR has been merged into BasicSR and this repo is a mirror of BasicSR.
https://github.com/xinntao/BasicSR
1.48k stars 320 forks source link

What can be reason why I get this error: TypeError: DCNv2Pack() takes no arguments ? #197

Open edgar2597 opened 3 years ago

edgar2597 commented 3 years ago

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

automatic0313 commented 2 years ago

I have encountered the same problem as you. Have you solved it?

BRIOTPatrick commented 2 years ago

model = EDVR_arch.EDVR(128, N_in, 8, 5, backRBs, 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)

in () 5 6 edvr_video(source_url, Path('/content/video.mp4'), data_mode, ----> 7 original_video_quality, 100, finetune_stage2) 8 9 # clear_output()

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