| 43 | return TF.resize(image, size=(new_h, new_w), interpolation=tf.InterpolationMode.BILINEAR) |
| 44 | |
| 45 | class DAVISDataset(Dataset): |
| 46 | def __init__(self, opt, training=True, shuffle=False): |
| 47 | self.opt = opt |
| 48 | self.training = training |
| 49 | self.shuffle = shuffle |
| 50 | self.split = "train" if training else "val" |
| 51 | |
| 52 | self.set_txt_path = osp.join(opt.root_path, "ImageSets", "2017", f"{self.split}.txt") |
| 53 | self.images_root = osp.join(opt.root_path, "JPEGImages", "Full-Resolution") |
| 54 | |
| 55 | with open(self.set_txt_path, "r") as f: |
| 56 | self.set_names = [line.strip() for line in f.readlines() if line.strip()] |
| 57 | if len(self.set_names) == 0: |
| 58 | raise RuntimeError(f"No set names found in {self.set_txt_path}") |
| 59 | |
| 60 | self.frame_infos = [] |
| 61 | self.all_frames = [] |
| 62 | self.set_len = [] # number of frames per set |
| 63 | |
| 64 | for set_idx, set_name in enumerate(self.set_names): |
| 65 | set_dir = osp.join(self.images_root, set_name) |
| 66 | image_files = glob.glob(osp.join(set_dir, "*.jpg")) |
| 67 | if not image_files: |
| 68 | print(f"Warning: No jpg images found in {set_dir}") |
| 69 | continue |
| 70 | |
| 71 | def extract_timestamp(file_path): |
| 72 | base = osp.basename(file_path) |
| 73 | num_str = osp.splitext(base)[0] |
| 74 | try: |
| 75 | return float(num_str) |
| 76 | except ValueError: |
| 77 | return 0.0 |
| 78 | |
| 79 | image_files = sorted(image_files, key=lambda x: extract_timestamp(x)) |
| 80 | |
| 81 | frames_for_set = [] |
| 82 | for idx, img_path in enumerate(image_files): |
| 83 | timestamp = extract_timestamp(img_path) |
| 84 | frames_for_set.append((img_path, timestamp)) |
| 85 | self.all_frames.append((set_idx, img_path, timestamp, idx)) |
| 86 | self.frame_infos.append(frames_for_set) |
| 87 | self.set_len.append(len(frames_for_set)) |
| 88 | |
| 89 | self.set_transform() |
| 90 | |
| 91 | def set_transform(self): |
| 92 | self.transform = tf.Compose([ |
| 93 | tf.ToTensor(), |
| 94 | ]) |
| 95 | |
| 96 | def get_predicted_depth_path(self, image_path): |
| 97 | pred_dir = image_path.replace("JPEGImages", "depthImages") |
| 98 | pred_dir = osp.dirname(pred_dir) |
| 99 | os.makedirs(pred_dir, exist_ok=True) |
| 100 | base_name = osp.splitext(osp.basename(image_path))[0] + "_pred.png" |
| 101 | return osp.join(pred_dir, base_name) |
| 102 | |