Method construct_sample

flow_poke/data.py:191–234

(sample: dict[str, bytes])

    def make_loader(self):
        def construct_sample(sample: dict[str, bytes]) -> dict[str, torch.Tensor | bool]:
            try:
                required_keys = ["video.mpg", "tracks.npy", "visibility.npy"]
                assert all(k in sample for k in required_keys), f"Expected keys {required_keys}, got {sample.keys()}"

                d = {
                    "tracks": torch.from_numpy(decode_npy(sample["tracks.npy"])),  # [t, n_t, 2]
                    "visibility": torch.from_numpy(decode_npy(sample["visibility.npy"])),  # [t, n_t]
                }

                sample_out = self.extract_training_sample(d)
                if not sample_out.get("valid", True):
                    return {"valid": False}
                sample_out = self.build_targets(sample_out)
                if not sample_out.get("valid", True):
                    return {"valid": False}

                i_f = sample_out.pop("i_frame")
                # Retrieve target frame from video
                with io.BytesIO(sample["video.mpg"]) as buf, av.open(buf) as container:
                    c_f = 0
                    target_frame = None
                    for packet in container.demux():
                        if not target_frame is None:
                            break
                        for frame in packet.decode():
                            if c_f == i_f:
                                target_frame = frame.to_ndarray(format="rgb24")
                                break
                            c_f += 1
                assert c_f == i_f, f"{i_f=}, {c_f=}"
                # Resize and normalize the target frame
                # TODO: augment -> explicit
                x: Float[torch.Tensor, "c h w"] = augment(
                    einops.rearrange(torch.from_numpy(target_frame).float() / 255, "h w c -> 1 c h w"), size=512
                )[
                    0
                ]  # [-1, 1]
                return sample_out | {
                    "x": x,
                }
            except Exception as e:
                print(f"Error while constructing sample : {e}")
                return {"valid": False}

        dataset = wds.DataPipeline(
            wds.ResampledShards(urls=([str(f) for f in Path(self.tar_base).rglob("*.tar")])),