<jupyter_start><jupyter_text>Check what processor Colab made avaiable to you:<jupyter_code>import torch
torch.cuda.get_device_name(0)<jupyter_output><empty_output><jupyter_text>Clone the project<jupyter_code>!git clone https://github.com/BurguerJohn/DAIN<jupyter_output><empty_output><jupyter_text>Install requirements<jupyter_code>!pip install torch==1.3.1
!pip install scipy==1.1.0<jupyter_output><empty_output><jupyter_text>Build ScriptsThis can take between 5 and 20 minutes.<jupyter_code>%%bash
cd /content/DAIN/my_package
sh build.sh
%%bash
cd /content/DAIN/PWCNet/correlation_package_pytorch1_0
sh build.sh<jupyter_output><empty_output><jupyter_text>Required models Download<jupyter_code>!mkdir /content/DAIN/MegaDepth/checkpoints/
!mkdir /content/DAIN/MegaDepth/checkpoints/test_local
%cd /content/DAIN/MegaDepth/checkpoints/test_local
!wget http://vllab1.ucmerced.edu/~wenbobao/DAIN/best_generalization_net_G.pth
%cd /content/DAIN/PWCNet
!wget http://vllab1.ucmerced.edu/~wenbobao/DAIN/pwc_net.pth.tar<jupyter_output><empty_output><jupyter_text>Optional model Download.This is the model used in Dain-App, only download it if you want to keep training on top of the existing model<jupyter_code>%cd /content/
!wget http://vllab1.ucmerced.edu/~wenbobao/DAIN/best.pth<jupyter_output><empty_output><jupyter_text>Tool to generate a DataSetPut all your videos on a folder and run this script to generate triplets [3 PNG's sequence] from all the videos inside that folder.<jupyter_code>#@markdown Folder path with all the videos
inputPath = "/content/input/" #@param{type:"string"}
#@markdown Path to generate the Database
outputPath = "/content/dataset_01/" #@param{type:"string"}
#@markdown The bigger the value, the more alike the frames need to be to be considered a triplet.
psnr = 25 #@param{type:"number"}
#@markdown The original model was trained with images of res 448X256, this bool turn all frames to that resolution.
resize = True #@param{type:"boolean"}

#Original dataset information:
#Vimeo-90K triplets dataset contains 91701 triplets extracted from 15k video clips.
#Each triplet is a short RGB video sequence that consists of 3 frames with fixed resolution 448x256


def calculate_psnr(img1, img2, max_value=255):
    """"Calculating peak signal-to-noise ratio (PSNR) between two images."""
    i1 = PIL.Image.open(img1).convert('RGB')
    i2 = PIL.Image.open(img2).convert('RGB')

    mse = np.mean((np.array(i1, dtype=np.float32) - np.array(i2, dtype=np.float32)) ** 2)
    if mse == 0:
        return 100
    return 20 * np.log10(max_value / (np.sqrt(mse)))


def IsDiffScenes(img1, img2, diff=25):
  return calculate_psnr(img1, img2) <= diff


import os
import shutil
from shutil import copyfile
import time
import PIL
import numpy as np


if not os.path.isdir(outputPath):
  os.makedirs(outputPath)

if not os.path.isdir(outputPath + "sequences"):
  os.makedirs(outputPath + "sequences")

counter = 1
for video in sorted(os.listdir(inputPath)):
  temp = "/content/.temp/"
  path = inputPath + video
  if os.path.isdir(temp):
    shutil.rmtree(temp)
    #time.sleep(2)
  os.makedirs(temp)

  rescale = ""
  if resize:
    rescale="scale=448:-1,pad=ceil(iw/2)*2:256,"
  !ffmpeg -i "{path}" -vf "{rescale}mpdecimate" -vsync 0 -qscale:v 1 -pix_fmt rgb24  "{temp}%10d.png"

  root =outputPath + "/sequences/" + str(counter).zfill(3) + "/"
  if not os.path.isdir(root):
    os.makedirs(root)


  triIndex = 1
  images = sorted(os.listdir(temp))
  for i in range(2, len(images)) :
    im1 = temp + images[i-2]
    im2 =  temp + images[i-1]
    im3 =  temp + images[i]
    if IsDiffScenes(im1, im2, psnr) == False and IsDiffScenes(im2, im3, psnr) == False:
      #print("Triplet Found")
      triPath = root + str(triIndex).zfill(5) + "/"
      if not os.path.isdir(triPath):
        os.makedirs(triPath)
      copyfile(im1, triPath + "im1.png")
      copyfile(im2, triPath + "im2.png")
      copyfile(im3, triPath + "im3.png")
      triIndex +=1

  counter += 1<jupyter_output><empty_output><jupyter_text>Tool to generate the training and testing set of a DataSet<jupyter_code>#@markdown The folder path of the generated Dataset
database = "/content/dataset_01/" #@param{type:"string"}
#@markdown Percentage of frames that will be to the testing group.
testingPercentage = 0.2 #@param{type:"slider", min:0.1, max:0.9, step:0.1}

import random

if os.path.isfile(database + "tri_testlist.txt"):
  os.remove(database + "tri_testlist.txt")
if os.path.isfile(database + "tri_trainlist.txt"):
  os.remove(database + "tri_trainlist.txt")

from glob import glob
folders = glob(database + "sequences/*/")

paths = []

for i in range(0, len(folders)):
  triplets = glob(folders[i] + "/*/")
  for ii in range(0, len(triplets)):
    innerF = os.path.basename((os.path.dirname(triplets[ii])))
    outF = os.path.basename((os.path.dirname(folders[i])))
    paths.append(outF + "/" + innerF)

def partitionRankings(rawRatings, testPercent):
    howManyNumbers = int(round(testPercent*len(rawRatings)))
    shuffled = rawRatings[:]
    random.shuffle(shuffled)
    return shuffled[howManyNumbers:], shuffled[:howManyNumbers]

training, test = partitionRankings(paths, testingPercentage)


trainTxt = database + "/tri_trainlist.txt"
trainFile = open(trainTxt, 'w')

testTxt = database + "/tri_testlist.txt"
testFile = open(testTxt, 'w')

for txt in training:
  trainFile.write(txt + "\n")
trainFile.close()

for txt in test:
  testFile.write(txt + "\n")
testFile.close()

print(training)
print(test)<jupyter_output><empty_output><jupyter_text>Train the model (This take a long time)<jupyter_code>%cd /content/DAIN
#@markdown Folder Path with the database
datapath_set = "/content/dataset_01" #@param{type:"string"}
#@markdown Empty if you want to start a new model from scratch.
#@markdown Set the path for the older model if you want to keep training it.
model_weigh = "" #@param{type:"string"}
#@markdown Size of the batch
batch_size = 3 #@param{type:"number"}

!CUDA_VISIBLE_DEVICES=0 python -W ignore train.py --datasetPath "{datapath_set}" --SAVED_MODEL "{model_weigh}" --batch_size {batch_size} --save_which 1 --lr 0.0005 --rectify_lr 0.0005 --flow_lr_coe 0.01 --occ_lr_coe 0.0 --filter_lr_coe 1.0 --ctx_lr_coe 1.0 --alpha 0.0 1.0 --patience 4 --factor 0.2<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup