Buckets:
| """Canonical semantic labels and categorical transforms for approved RGB tiles.""" | |
| from dataclasses import dataclass, replace | |
| from pathlib import Path | |
| import json | |
| import numpy as np | |
| from PIL import Image | |
| from load_dataset import dataset, load_sample | |
| CANONICAL_IDS = (1, 2, 12, 13, 14, 15, 16) | |
| CLASS_NAMES = ("building", "road", "water", "vegetation", "agricultural_area", "bare_ground", "developed_area") | |
| def canonical_to_head(mask): | |
| mask = np.asarray(mask) | |
| if mask.ndim != 2 or not np.issubdtype(mask.dtype, np.integer) or not np.isin(mask, (*CANONICAL_IDS, 255)).all(): | |
| raise ValueError("expected supported integer canonical labels or ignore255") | |
| result = np.full(mask.shape, 255, np.uint8) | |
| for index, canonical in enumerate(CANONICAL_IDS): | |
| result[mask == canonical] = index | |
| return result | |
| def head_to_canonical(mask): | |
| mask = np.asarray(mask) | |
| if mask.ndim != 2 or not np.issubdtype(mask.dtype, np.integer) or not np.isin(mask, (*range(7), 255)).all(): | |
| raise ValueError("expected contiguous semantic labels or ignore255") | |
| result = np.full(mask.shape, 255, np.uint8) | |
| for index, canonical in enumerate(CANONICAL_IDS): | |
| result[mask == index] = canonical | |
| return result | |
| class SegmentationExample: | |
| image: np.ndarray | |
| mask: np.ndarray | |
| head_mask: np.ndarray | |
| valid_mask: np.ndarray | |
| metadata: dict | |
| def convert_segmentation_sample(sample): | |
| if sample["metadata"]["task"] != "segmentation": | |
| raise ValueError("semantic sample required") | |
| mask = np.array(sample["mask"], copy=True) | |
| canonical_to_head(mask) # validate before excluding invalid pixels | |
| valid = np.asarray(sample["valid_mask"]) | |
| if valid.dtype != bool or valid.shape != mask.shape or mask.shape != sample["image"].shape[:2]: | |
| raise ValueError("boolean validity mask must match image and labels") | |
| valid = valid & (mask != 255) | |
| mask[~valid] = 255 | |
| return SegmentationExample(sample["image"], mask, canonical_to_head(mask), valid, dict(sample["metadata"])) | |
| def load_segmentation_tile(root: Path, sample: dict): | |
| approved = dataset(root, "segmentation", sample["split"], license_scope="all") | |
| if {r["id"]: r for r in approved.samples}.get(sample["id"]) != sample: | |
| raise ValueError("sample differs from approved manifest") | |
| return convert_segmentation_sample(load_sample(root, sample, approved.expected_hashes)) | |
| def transform_example(example, *, horizontal=False, vertical=False, rotations=0, size=None): | |
| image, mask = example.image, example.mask | |
| if horizontal: | |
| image, mask = image[:, ::-1], mask[:, ::-1] | |
| if vertical: | |
| image, mask = image[::-1], mask[::-1] | |
| image, mask = np.rot90(image, rotations % 4), np.rot90(mask, rotations % 4) | |
| if size is not None: | |
| if len(size) != 2 or any(type(s) is not int or s <= 0 for s in size): | |
| raise ValueError("positive integer (height,width) required") | |
| image = np.array(Image.fromarray(image).resize((size[1], size[0]), Image.Resampling.BILINEAR)) | |
| mask = np.array(Image.fromarray(mask).resize((size[1], size[0]), Image.Resampling.NEAREST)) | |
| return replace(example, image=np.ascontiguousarray(image), mask=np.ascontiguousarray(mask), | |
| head_mask=canonical_to_head(mask), valid_mask=np.ascontiguousarray(mask != 255)) | |
| def segmentation_batch(examples): | |
| import torch | |
| return {"img": torch.from_numpy(np.stack([e.image.transpose(2, 0, 1) for e in examples])).float() / 255, | |
| "semantic_mask": torch.from_numpy(np.stack([e.head_mask for e in examples])).long(), | |
| "valid_mask": torch.from_numpy(np.stack([e.valid_mask for e in examples]))} | |
| def build_segmentation_cache(root: Path, output: Path): | |
| repo = Path(__file__).resolve().parents[1] | |
| output = output.resolve() | |
| if not output.is_relative_to(repo / "training"): | |
| raise ValueError("cache must be under training/") | |
| data = dataset(root, "segmentation", "train", license_scope="all") | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| with output.open("x", encoding="utf-8") as handle: | |
| json.dump({"corpus_content_sha256": json.loads((Path(root) / "stats.json").read_text())["validation_content_sha256"], | |
| "canonical_ids": CANONICAL_IDS, "samples": data.samples}, handle) | |
| return output | |
Xet Storage Details
- Size:
- 4.38 kB
- Xet hash:
- 300006117212f7c8bd7d674288b317eb2afe359de356fc47587a8fe00ce44abc
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.