Spaces:
Configuration error
Configuration error
File size: 1,404 Bytes
d01f62c 9e837cb d01f62c 9e837cb d01f62c 9e837cb d01f62c 9e837cb 0580cf1 9e837cb d01f62c 9e837cb d01f62c 9e837cb 0580cf1 9e837cb d01f62c 9e837cb | 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 | import os
from os import path
import json
from inference.data.video_reader import VideoReader_221128_TransColorization
class DAVISTestDataset_221128_TransColorization_batch:
def __init__(self, data_root, imset='2017/val.txt', size=-1, args=None):
self.image_dir = data_root
self.mask_dir = imset
self.size_dir = data_root
self.size = size
self.vid_list = [clip_name for clip_name in sorted(os.listdir(data_root)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
self.ref_img_list = [clip_name for clip_name in sorted(os.listdir(imset)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
self.args = args
# print(lst, len(lst), self.vid_list, self.vid_list_DAVIS2016, path.join(data_root, 'ImageSets', imset));assert 1==0
def get_datasets(self):
for video in self.vid_list:
if video not in self.ref_img_list:
continue
# print(self.image_dir, video, path.join(self.image_dir, video));assert 1==0
yield VideoReader_221128_TransColorization(video,
path.join(self.image_dir, video),
path.join(self.mask_dir, video),
size=self.size,
size_dir=path.join(self.size_dir, video),
args=self.args
)
def __len__(self):
return len(self.vid_list)
|