File size: 2,968 Bytes
872b0a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import os
import sys
import re
import numpy as np
from PIL import Image
from ..core.data import Data
from ..util import tryremove
from shutil import copyfile, rmtree
from urllib.request import urlretrieve
class ChairsData(Data):
URL = 'http://lmb.informatik.uni-freiburg.de/data/FlyingChairs/FlyingChairs.zip'
TRAIN_VAL_URL = 'http://lmb.informatik.uni-freiburg.de/resources/datasets/FlyingChairs/FlyingChairs_train_val.txt'
dirs = ['flying_chairs']
def __init__(self, data_dir, stat_log_dir=None,
development=True, fast_dir=None):
super().__init__(data_dir, stat_log_dir,
development=development,
fast_dir=fast_dir)
def _fetch_if_missing(self):
local_path = os.path.join(self.data_dir, 'flying_chairs')
train_val_path = os.path.join(local_path, 'FlyingChairs_train_val.txt')
if not os.path.isdir(local_path):
did_download = True
self._download_and_extract(self.URL, local_path)
urlretrieve(self.TRAIN_VAL_URL, train_val_path)
else:
did_download = False
data_path = os.path.join(local_path, 'FlyingChairs_release', 'data')
os.makedirs(os.path.join(local_path, 'image'), exist_ok=True)
os.makedirs(os.path.join(local_path, 'flow'), exist_ok=True)
os.makedirs(os.path.join(local_path, 'test_image'), exist_ok=True)
if os.path.isdir(data_path):
print('>> converting chairs data to .png')
train_val_repeated = []
train_val = []
with open(train_val_path) as f:
for line in f:
training = int(line.strip()) == 1
train_val_repeated.extend([training, training])
train_val.extend([training])
# Convert .ppm to .png and split data into image and flow directory
im_files = [f for f in os.listdir(data_path) if
re.match(r'[0-9]+.*\.ppm', f)]
im_files.sort()
flow_files = [f for f in os.listdir(data_path) if
re.match(r'[0-9]+.*\.flo', f)]
flow_files.sort()
for t, f in zip(train_val_repeated, im_files):
name, ext = os.path.splitext(f)
path = os.path.join(data_path, f)
im = Image.open(path)
folder = 'image' if t else 'test_image'
im.save(os.path.join(local_path, folder, name + '.png'),
'PNG')
for t, f in zip(train_val, flow_files):
path = os.path.join(data_path, f)
if not t:
copyfile(path, os.path.join(local_path, 'flow', f))
if did_download:
rmtree(data_path)
print('>> processed chairs data')
def get_raw_dirs(self):
return [os.path.join(self.current_dir, 'flying_chairs', 'image')]
|