Read a segmentation mask ### Parameters: - `path: Union[str, os.PathLike, IO]` The file path or file object to read from. ### Returns: - `Tuple[np.ndarray, Dict[str, int]]` A tuple containing: - `mask`: uint8 or uint16 numpy.ndarray of shape (H, W).
(path: Union[str, os.PathLike, IO])
| 155 | |
| 156 | |
| 157 | def read_segmentation(path: Union[str, os.PathLike, IO]) -> Tuple[np.ndarray, Dict[str, int]]: |
| 158 | """ |
| 159 | Read a segmentation mask |
| 160 | ### Parameters: |
| 161 | - `path: Union[str, os.PathLike, IO]` |
| 162 | The file path or file object to read from. |
| 163 | ### Returns: |
| 164 | - `Tuple[np.ndarray, Dict[str, int]]` |
| 165 | A tuple containing: |
| 166 | - `mask`: uint8 or uint16 numpy.ndarray of shape (H, W). |
| 167 | - `labels`: Dict[str, int]. The label mapping, a dictionary of {label_name: label_id}. |
| 168 | """ |
| 169 | if isinstance(path, (str, os.PathLike)): |
| 170 | data = Path(path).read_bytes() |
| 171 | else: |
| 172 | data = path.read() |
| 173 | pil_image = Image.open(io.BytesIO(data)) |
| 174 | labels = json.loads(pil_image.info['labels']) if 'labels' in pil_image.info else None |
| 175 | mask = np.array(pil_image) |
| 176 | return mask, labels |
| 177 | |
| 178 | |
| 179 | def write_segmentation(path: Union[str, os.PathLike, IO], mask: np.ndarray, labels: Dict[str, int] = None, compression_level: int = 7): |