| 35 | # --------------------------------------------------------------------------------------------------------------------- |
| 36 | |
| 37 | class Example: |
| 38 | |
| 39 | def __init__(self, img_path: str, background_queries: list, object_queries: list, object_pokes: list): |
| 40 | self.img_path = img_path |
| 41 | self.background_queries = background_queries |
| 42 | self.background_pokes = [(0.0, 0.0)]*len(background_queries) |
| 43 | self.object_queries = object_queries |
| 44 | self.object_pokes = object_pokes |
| 45 | |
| 46 | def get_image(self): |
| 47 | return Image.open(self.img_path).convert("RGB") |
| 48 | |
| 49 | def get_queries(self) -> list[tuple[float, float]]: |
| 50 | return self.background_queries + self.object_queries |
| 51 | |
| 52 | def get_pokes(self, mode='fully_poked') -> list[tuple[float, float]] | None: |
| 53 | if mode == 'fully_poked': |
| 54 | return self.background_pokes + self.object_pokes |
| 55 | if mode == "first_query": |
| 56 | return self.background_pokes + self.object_pokes[:1] |
| 57 | elif mode == 'unpoked': |
| 58 | return None |
| 59 | elif mode == 'background_fixed': |
| 60 | return self.background_pokes + [(np.nan, np.nan)]*len(self.object_queries) |
| 61 | elif mode == 'object_coherence': |
| 62 | return self.background_pokes + self.object_pokes[:1] + [(np.nan, np.nan)]*len(self.object_pokes[1:]) |
| 63 | else: |
| 64 | raise NotImplementedError(f"Mode {mode} not implemented") |
| 65 | |
| 66 | def get_pokes_queries(self, mode='fully_poked'): |
| 67 | if mode == 'fully_poked': |
| 68 | return self.background_pokes + self.object_pokes, self.get_queries() |
| 69 | if mode == "first_query": |
| 70 | return self.background_pokes + self.object_pokes[:1], self.background_queries + self.object_queries[:1] |
| 71 | elif mode == 'unpoked': |
| 72 | return None, self.get_queries() |
| 73 | elif mode == 'background_fixed': |
| 74 | return self.background_pokes + [(np.nan, np.nan)]*len(self.object_queries), self.get_queries() |
| 75 | elif mode == 'object_coherence': |
| 76 | return self.background_pokes + self.object_pokes[:1] + [(np.nan, np.nan)]*len(self.object_pokes[1:]), self.get_queries() |
| 77 | else: |
| 78 | raise NotImplementedError(f"Mode {mode} not implemented") |
| 79 | |
| 80 | EXAMPLES = { |
| 81 | "falling_ball": Example( |