diff --git a/codes/scripts/extract_temporal_squares.py b/codes/scripts/extract_temporal_squares.py
index e69de29b..bb74f7ed 100644
--- a/codes/scripts/extract_temporal_squares.py
+++ b/codes/scripts/extract_temporal_squares.py
@@ -0,0 +1,201 @@
+"""A multi-thread tool to crop large images to sub-images for faster IO."""
+import os
+import os.path as osp
+import shutil
+import subprocess
+from time import sleep
+
+import munch
+import numpy as np
+import cv2
+import torchvision
+from PIL import Image
+import data.util as data_util # noqa: E402
+import torch.utils.data as data
+from tqdm import tqdm
+import torch
+import random
+
+from models.flownet2.networks.resample2d_package.resample2d import Resample2d
+from models.optical_flow.PWCNet import pwc_dc_net
+
+
+def main():
+ opt = {}
+ opt['n_thread'] = 0
+ opt['compression_level'] = 95 # JPEG compression quality rating.
+ opt['dest'] = 'file'
+ opt['input_folder'] = 'D:\\dlas\\codes\\scripts\\test'
+ opt['save_folder'] = 'D:\\dlas\\codes\\scripts\\test_out'
+ opt['imgsize'] = 256
+ opt['bottom_crop'] = .1
+ opt['keep_folder'] = False
+
+ save_folder = opt['save_folder']
+ if not osp.exists(save_folder):
+ os.makedirs(save_folder)
+ print('mkdir [{:s}] ...'.format(save_folder))
+
+ go(opt)
+
+
+def is_video(filename):
+ return any(filename.endswith(extension) for extension in ['.mp4', '.MP4', '.avi', '.AVI', '.mkv', '.MKV', '.wmv', '.WMV'])
+
+
+def get_videos_in_path(path):
+ assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+ videos = []
+ for dirpath, _, fnames in sorted(os.walk(path)):
+ for fname in sorted(fnames):
+ if is_video(fname):
+ videos.append(os.path.join(dirpath, fname))
+ return videos
+
+def get_time_for_secs(secs):
+ mins = int(secs / 60)
+ hours = int(mins / 60)
+ secs = secs - (mins * 60) - (hours * 3600)
+ mins = mins % 60
+ return '%02d:%02d:%06.3f' % (hours, mins, secs)
+
+class VideoClipDataset(data.Dataset):
+ def __init__(self, opt):
+ self.opt = opt
+ input_folder = opt['input_folder']
+ self.videos = get_videos_in_path(input_folder)
+ print("Found %i videos" % (len(self.videos),))
+
+ def __getitem__(self, index):
+ return self.get(index)
+
+ def extract_n_frames(self, video_file, dest, time_seconds, n):
+ ffmpeg_args = ['ffmpeg', '-y', '-ss', get_time_for_secs(time_seconds), '-i', video_file, '-vframes', str(n), f'{dest}/%d.jpg']
+ process = subprocess.Popen(ffmpeg_args, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
+ process.wait()
+
+ def get_video_length(self, video_file):
+ result = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of",
+ "default=noprint_wrappers=1", video_file],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT)
+ return float(result.stdout.decode('utf-8').strip().replace("duration=", ""))
+
+ def get_image_tensor(self, path):
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+
+ # Access exceptions happen regularly, probably due to the subprocess not fully terminating.
+ for tries in range(5):
+ try:
+ os.remove(path)
+ break
+ except:
+ if tries >= 4:
+ assert False
+ else:
+ sleep(.1)
+
+ assert img is not None
+ assert len(img.shape) > 2
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
+ # Crop off excess so image dimensions are a multiple of 64.
+ h, w, _ = img.shape
+ img = img[:(h//64)*64,:(w//64)*64,:]
+ return torch.from_numpy(np.ascontiguousarray(np.transpose(img, (2, 0, 1)))).float()
+
+ def get(self, index):
+ path = self.videos[index]
+ out_folder = self.opt['save_folder']
+
+ vid_len = int(self.get_video_length(path))
+ start = 2
+ img_runs = []
+ while start < vid_len:
+ frames_out = os.path.join(out_folder, f'{index}_{start}')
+ os.makedirs(frames_out, exist_ok=False)
+ n = random.randint(5, 30)
+ self.extract_n_frames(path, frames_out, start, n)
+ frames = data_util.get_image_paths('img', frames_out)[0]
+ assert len(frames) == n
+ img_runs.append(([self.get_image_tensor(frame) for frame in frames], frames_out))
+ start += random.randint(2,5)
+
+ return img_runs
+
+ def __len__(self):
+ return len(self.videos)
+
+
+def compute_flow_and_cleanup(flownet, runs):
+ resampler = Resample2d().cuda()
+ for run in runs:
+ run, path = run
+ consolidated_flows = None
+ a = run[0].unsqueeze(0).cuda()
+ img = a
+ dbg = a.clone()
+ for i in range(1,len(run)):
+ img2 = run[i].unsqueeze(0).cuda()
+ flow = flownet(torch.cat([img2, img], dim=1))
+ flow = torch.nn.functional.interpolate(flow, size=img.shape[2:], mode='bilinear')
+ if consolidated_flows is None:
+ consolidated_flows = flow
+ else:
+ consolidated_flows = resampler(flow, -consolidated_flows) + consolidated_flows
+ img = img2
+ dbg = resampler(dbg, flow)
+ torchvision.utils.save_image(dbg, os.path.join(path, "debug.jpg"))
+ consolidated_flows = torch.clamp(consolidated_flows / 255, -.5, .5)
+ b = run[-1].unsqueeze(0).cuda()
+ _, _, h, w = a.shape
+ direct_flows = torch.nn.functional.interpolate(torch.clamp(flownet(torch.cat([a, b], dim=1).float()) / 255, -.5, .5), size=img.shape[2:], mode='bilinear')
+ # TODO: Reshape image here.
+ '''
+ # Perform explicit crops first. These are generally used to get rid of watermarks so we dont even want to
+ # consider these areas of the image.
+ if 'bottom_crop' in self.opt.keys() and self.opt['bottom_crop'] > 0:
+ bc = self.opt['bottom_crop']
+ if bc > 0 and bc < 1:
+ bc = int(bc * img.shape[0])
+ img = img[:-bc, :, :]
+
+ h, w, c = img.shape
+ assert min(h,w) >= self.opt['imgsize']
+
+ # We must convert the image into a square.
+ dim = min(h, w)
+ # Crop the image so that only the center is left, since this is often the most salient part of the image.
+ img = img[(h - dim) // 2:dim + (h - dim) // 2, (w - dim) // 2:dim + (w - dim) // 2, :]
+ img = cv2.resize(img, (self.opt['imgsize'], self.opt['imgsize']), interpolation=cv2.INTER_AREA)
+ '''
+
+ torchvision.utils.save_image(a, os.path.join(path, "a.jpg"))
+ torchvision.utils.save_image(b, os.path.join(path, "b.jpg"))
+ torch.save(consolidated_flows * 255, os.path.join(path, "consolidated_flow.pt"))
+ torchvision.utils.save_image(torch.cat([consolidated_flows + .5, torch.zeros((1, 1, h, w), device='cuda')], dim=1), os.path.join(path, "consolidated_flow.png"))
+
+ # For debugging
+ torchvision.utils.save_image(resampler(a, consolidated_flows * 255), os.path.join(path, "b_flowed.jpg"))
+ torchvision.utils.save_image(resampler(b, -consolidated_flows * 255), os.path.join(path, "a_flowed.jpg"))
+ torchvision.utils.save_image(resampler(b, direct_flows * 255), os.path.join(path, "a_flowed_nonconsolidated.jpg"))
+ torchvision.utils.save_image(torch.cat([direct_flows + .5, torch.zeros((1, 1, h, w), device='cuda')], dim=1), os.path.join(path, "direct_flow.png"))
+
+
+def identity(x):
+ return x
+
+
+def go(opt):
+ flownet = pwc_dc_net('../experiments/pwc_humanflow.pth')
+ flownet.eval()
+ flownet = flownet.cuda()
+
+ dataset = VideoClipDataset(opt)
+ dataloader = data.DataLoader(dataset, num_workers=opt['n_thread'], collate_fn=identity)
+ with torch.no_grad():
+ for batch in tqdm(dataloader):
+ compute_flow_and_cleanup(flownet, batch[0])
+
+
+if __name__ == '__main__':
+ main()