| 85 | |
| 86 | |
| 87 | class Assembly: |
| 88 | def __init__(self, size): |
| 89 | self.data = np.full((size, 4), np.nan) |
| 90 | self.confidence = 0 # 0 by default, overwritten otherwise with `add_joint` |
| 91 | self._affinity = 0 |
| 92 | self._links = [] |
| 93 | self._visible = set() |
| 94 | self._idx = set() |
| 95 | self._dict = dict() |
| 96 | |
| 97 | def __len__(self): |
| 98 | return len(self._visible) |
| 99 | |
| 100 | def __contains__(self, assembly): |
| 101 | return bool(self._visible.intersection(assembly._visible)) |
| 102 | |
| 103 | def __add__(self, other): |
| 104 | if other in self: |
| 105 | raise ValueError("Assemblies contain shared joints.") |
| 106 | |
| 107 | assembly = Assembly(self.data.shape[0]) |
| 108 | for link in self._links + other._links: |
| 109 | assembly.add_link(link) |
| 110 | return assembly |
| 111 | |
| 112 | @classmethod |
| 113 | def from_array(cls, array): |
| 114 | n_bpts, n_cols = array.shape |
| 115 | |
| 116 | # if a single coordinate is NaN for a bodypart, set all to NaN |
| 117 | array[np.isnan(array).any(axis=-1)] = np.nan |
| 118 | |
| 119 | ass = cls(size=n_bpts) |
| 120 | ass.data[:, :n_cols] = array |
| 121 | visible = np.flatnonzero(~np.isnan(array).any(axis=1)) |
| 122 | if n_cols < 3: # Only xy coordinates are being set |
| 123 | ass.data[visible, 2] = 1 # Set detection confidence to 1 |
| 124 | ass._visible.update(visible) |
| 125 | return ass |
| 126 | |
| 127 | @property |
| 128 | def xy(self): |
| 129 | return self.data[:, :2] |
| 130 | |
| 131 | @property |
| 132 | def extent(self): |
| 133 | bbox = np.empty(4) |
| 134 | bbox[:2] = np.nanmin(self.xy, axis=0) |
| 135 | bbox[2:] = np.nanmax(self.xy, axis=0) |
| 136 | return bbox |
| 137 | |
| 138 | @property |
| 139 | def area(self): |
| 140 | x1, y1, x2, y2 = self.extent |
| 141 | return (x2 - x1) * (y2 - y1) |
| 142 | |
| 143 | @property |
| 144 | def confidence(self): |
no outgoing calls
no test coverage detected