MCPcopy Create free account
hub / github.com/DSL-Lab/StreamSplat / __getitem__

Method __getitem__

datasets/provider_davis.py:100–177  ·  view source on GitHub ↗
(self, index)

Source from the content-addressed store, hash-verified

98 return len(self.ret_frames)
99
100 def __getitem__(self, index):
101 set_idx, image_path, timestamp, frame_idx = self.ret_frames[index]
102
103 if not osp.exists(image_path):
104 raise FileNotFoundError(f"Image not found: {image_path}")
105
106 main_image = Image.open(image_path).convert("RGB")
107 depth_path = self.get_predicted_depth_path(image_path)
108 depth_image = Image.open(depth_path)
109 main_image = self.transform(main_image)
110 depth_image = self.depth_transform(depth_image)
111
112 mask_path = image_path.replace(self.images_root, self.mask_root).replace(".jpg", ".png")
113 if not osp.exists(mask_path):
114 raise FileNotFoundError(f"Mask not found: {mask_path}")
115 mask_image = Image.open(mask_path)
116 mask_image = self.transform_mask(mask_image).unsqueeze(0)
117
118 frames = [main_image]
119 depths = [depth_image]
120 masks = [mask_image]
121 timestamps_list = [0.0]
122
123 seq_length = self.set_len[set_idx]
124 current_idx = frame_idx
125
126 if self.opt.output_frames > 1 and seq_length > 1:
127 if not self.shuffle:
128 if current_idx + self.nearby_range >= seq_length:
129 interval = (seq_length - current_idx) // (self.opt.output_frames - 1)
130 else:
131 interval = self.nearby_range // (self.opt.output_frames - 1)
132 assert interval > 0, "Interval must be greater than 0"
133 offsets = [i * interval for i in range(1, self.opt.output_frames)]
134 else:
135 if current_idx + self.nearby_range >= seq_length:
136 offsets = random.sample(range(1, seq_length - current_idx), self.opt.output_frames - 1)
137 else:
138 offsets = random.sample(range(1, self.nearby_range + 1), self.opt.output_frames - 1)
139
140 offsets.sort()
141
142 for offset in offsets:
143 pair_idx = current_idx + offset
144 img_file_pair, ts_pair = self.frame_infos[set_idx][pair_idx]
145 if not osp.exists(img_file_pair):
146 print(f"Warning: Missing file {img_file_pair}")
147 continue
148 pair_image = Image.open(img_file_pair).convert("RGB")
149 pair_depth_path = self.get_predicted_depth_path(img_file_pair)
150 pair_depth_image = Image.open(pair_depth_path)
151 pair_image = self.transform(pair_image)
152 pair_depth_image = self.transform(pair_depth_image)
153
154 pair_mask_path = img_file_pair.replace(self.images_root, self.mask_root).replace(".jpg", ".png")
155 if not osp.exists(pair_mask_path):
156 print(f"Warning: Missing file {pair_mask_path}")
157 continue

Callers

nothing calls this directly

Calls 3

transform_maskMethod · 0.95
sampleMethod · 0.80

Tested by

no test coverage detected