Hsyshasuzh/longmemory-geo / project /training /segmentation_data.py
Hsyshasuzh's picture
download
raw
4.38 kB
"""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
@dataclass(frozen=True)
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.