Spaces:
Runtime error
Runtime error
File size: 4,170 Bytes
f239efc |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import logging
import os
import json
import random
from torch.utils.data import Dataset
import time
from dataset.utils import load_image_from_path
try:
from petrel_client.client import Client
has_client = True
except ImportError:
has_client = False
logger = logging.getLogger(__name__)
class ImageVideoBaseDataset(Dataset):
"""Base class that implements the image and video loading methods"""
media_type = "video"
def __init__(self):
assert self.media_type in ["image", "video", "only_video"]
self.data_root = None
self.anno_list = (
None # list(dict), each dict contains {"image": str, # image or video path}
)
self.transform = None
self.video_reader = None
self.num_tries = None
self.client = None
if has_client:
self.client = Client('~/petreloss.conf')
def __getitem__(self, index):
raise NotImplementedError
def __len__(self):
raise NotImplementedError
def get_anno(self, index):
"""obtain the annotation for one media (video or image)
Args:
index (int): The media index.
Returns: dict.
- "image": the filename, video also use "image".
- "caption": The caption for this file.
"""
anno = self.anno_list[index]
if self.data_root is not None:
anno["image"] = os.path.join(self.data_root, anno["image"])
return anno
def load_and_transform_media_data(self, index, data_path):
if self.media_type == "image":
return self.load_and_transform_media_data_image(index, data_path, clip_transform=self.clip_transform)
else:
return self.load_and_transform_media_data_video(index, data_path, clip_transform=self.clip_transform)
def load_and_transform_media_data_image(self, index, data_path, clip_transform=False):
image = load_image_from_path(data_path, client=self.client)
if not clip_transform:
image = self.transform(image)
return image, index
def load_and_transform_media_data_video(self, index, data_path, return_fps=False, clip=None, clip_transform=False):
for _ in range(self.num_tries):
try:
max_num_frames = self.max_num_frames if hasattr(self, "max_num_frames") else -1
if "webvid" in data_path:
hdfs_dir="hdfs://harunava/home/byte_ailab_us_cvg/user/weimin.wang/videogen_data/webvid_data/10M_full_train"
video_name = os.path.basename(data_path)
video_id, extension = os.path.splitext(video_name)
ind_file = os.path.join(hdfs_dir, self.keys_indexfile[video_id])
frames, frame_indices, fps = self.video_reader(ind_file, video_id, self.num_frames, self.sample_type,
max_num_frames=max_num_frames, client=self.client, clip=clip)
else:
frames, frame_indices, fps = self.video_reader(
data_path, self.num_frames, self.sample_type,
max_num_frames=max_num_frames, client=self.client, clip=clip
)
except Exception as e:
logger.warning(
f"Caught exception {e} when loading video {data_path}, "
f"randomly sample a new video as replacement"
)
index = random.randint(0, len(self) - 1)
ann = self.get_anno(index)
data_path = ann["image"]
continue
# shared aug for video frames
if not clip_transform:
frames = self.transform(frames)
if return_fps:
sec = [str(round(f / fps, 1)) for f in frame_indices]
return frames, index, sec
else:
return frames, index
else:
raise RuntimeError(
f"Failed to fetch video after {self.num_tries} tries. "
f"This might indicate that you have many corrupted videos."
)
|