| 9 | import numpy as np |
| 10 | |
| 11 | class VOSDataset(Dataset): |
| 12 | def __init__(self, opt, training=True, shuffle=False, nearby_range=3): |
| 13 | self.opt = opt |
| 14 | self.training = training |
| 15 | self.shuffle = shuffle |
| 16 | if nearby_range < opt.output_frames - 1: |
| 17 | print(f"Warning: nearby_range is less than output_frames - 1 on vos, nearby_range: {nearby_range}, output_frames: {opt.output_frames}, change nearby_range to {opt.output_frames - 1}") |
| 18 | nearby_range = opt.output_frames - 1 |
| 19 | self.nearby_range = nearby_range |
| 20 | |
| 21 | self.split = "train" if training else "val" |
| 22 | |
| 23 | base = osp.join(opt.root_path, "train" if training else "valid") |
| 24 | self.images_root = osp.join(base, "JPEGImages") |
| 25 | self.mask_root = osp.join(base, "Annotations") if training else None |
| 26 | |
| 27 | self.set_names = [ |
| 28 | d for d in os.listdir(self.images_root) |
| 29 | if osp.isdir(osp.join(self.images_root, d)) |
| 30 | ] |
| 31 | |
| 32 | self.frame_infos = [] |
| 33 | self._all_frames = [] |
| 34 | self.set_len = [] |
| 35 | self.ret_frames = [] |
| 36 | |
| 37 | for set_idx, set_name in enumerate(self.set_names): |
| 38 | set_dir = osp.join(self.images_root, set_name) |
| 39 | image_files = glob.glob(osp.join(set_dir, "*.jpg")) |
| 40 | if not image_files: |
| 41 | print(f"Warning: No jpg images found in {set_dir}") |
| 42 | continue |
| 43 | |
| 44 | def extract_timestamp(file_path): |
| 45 | base = osp.basename(file_path) |
| 46 | num_str = osp.splitext(base)[0] |
| 47 | try: |
| 48 | return float(num_str) |
| 49 | except ValueError: |
| 50 | return 0.0 |
| 51 | |
| 52 | image_files = sorted(image_files, key=lambda x: extract_timestamp(x)) |
| 53 | |
| 54 | frames_for_set = [] |
| 55 | for idx, img_path in enumerate(image_files): |
| 56 | timestamp = extract_timestamp(img_path) |
| 57 | frames_for_set.append((img_path, timestamp)) |
| 58 | self._all_frames.append((set_idx, img_path, timestamp, idx)) |
| 59 | |
| 60 | has_enough_frames = True |
| 61 | if self.opt.output_frames > 1 and idx + self.opt.output_frames > len(image_files): |
| 62 | has_enough_frames = False |
| 63 | if has_enough_frames: |
| 64 | self.ret_frames.append((set_idx, img_path, timestamp, idx)) |
| 65 | |
| 66 | self.frame_infos.append(frames_for_set) |
| 67 | self.set_len.append(len(frames_for_set)) |
| 68 | |