Spaces:
Running on Zero
Running on Zero
File size: 4,238 Bytes
543832d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Pure image-processing helpers for the street-scene Space."""
from __future__ import annotations
import csv
from pathlib import Path
from typing import Iterable
import numpy as np
from PIL import Image, ImageOps
MAX_OUTPUT_SIDE = 2048
# Official Cityscapes train-ID palette.
CITYSCAPES_PALETTE: dict[str, tuple[int, int, int]] = {
"road": (128, 64, 128),
"sidewalk": (244, 35, 232),
"building": (70, 70, 70),
"wall": (102, 102, 156),
"fence": (190, 153, 153),
"pole": (153, 153, 153),
"traffic light": (250, 170, 30),
"traffic sign": (220, 220, 0),
"vegetation": (107, 142, 35),
"terrain": (152, 251, 152),
"sky": (70, 130, 180),
"person": (220, 20, 60),
"rider": (255, 0, 0),
"car": (0, 0, 142),
"truck": (0, 0, 70),
"bus": (0, 60, 100),
"train": (0, 80, 100),
"motorcycle": (0, 0, 230),
"bicycle": (119, 11, 32),
}
def _normalise_label(label: str) -> str:
return label.lower().replace("_", " ").strip()
def _fallback_color(class_id: int) -> tuple[int, int, int]:
"""Return a deterministic, visually distinct color for an unknown class."""
return (
int((37 * class_id + 71) % 205 + 25),
int((67 * class_id + 29) % 205 + 25),
int((97 * class_id + 11) % 205 + 25),
)
def class_color(class_id: int, label: str) -> tuple[int, int, int]:
return CITYSCAPES_PALETTE.get(_normalise_label(label), _fallback_color(class_id))
def resize_for_output(image: Image.Image, max_side: int = MAX_OUTPUT_SIDE) -> Image.Image:
"""Bound output resolution so large phone photos do not exhaust Space memory."""
image = ImageOps.exif_transpose(image).convert("RGB")
width, height = image.size
longest = max(width, height)
if longest <= max_side:
return image
scale = max_side / longest
size = (max(1, round(width * scale)), max(1, round(height * scale)))
resampling = getattr(Image, "Resampling", Image)
return image.resize(size, resampling.LANCZOS)
def render_segmentation(
image: Image.Image,
class_map: np.ndarray,
id2label: dict[int, str],
opacity: float,
) -> tuple[Image.Image, Image.Image]:
"""Create a Cityscapes-color mask and an overlay with white boundaries."""
height, width = class_map.shape
color_array = np.zeros((height, width, 3), dtype=np.uint8)
for class_id in np.unique(class_map):
label = id2label.get(int(class_id), f"class_{int(class_id)}")
color_array[class_map == class_id] = class_color(int(class_id), label)
base_array = np.asarray(image, dtype=np.float32)
overlay_array = (
base_array * (1.0 - opacity) + color_array.astype(np.float32) * opacity
).astype(np.uint8)
boundaries = np.zeros((height, width), dtype=bool)
boundaries[1:, :] |= class_map[1:, :] != class_map[:-1, :]
boundaries[:, 1:] |= class_map[:, 1:] != class_map[:, :-1]
overlay_array[boundaries] = (255, 255, 255)
return Image.fromarray(overlay_array), Image.fromarray(color_array)
def build_class_table(
class_map: np.ndarray,
id2label: dict[int, str],
min_share_percent: float,
) -> list[list[object]]:
"""Summarise class coverage in descending order."""
class_ids, counts = np.unique(class_map, return_counts=True)
total_pixels = int(class_map.size)
rows: list[list[object]] = []
for class_id, count in zip(class_ids, counts):
share = 100.0 * int(count) / total_pixels
if share < min_share_percent:
continue
label = id2label.get(int(class_id), f"class_{int(class_id)}")
color = class_color(int(class_id), label)
rows.append(
[
int(class_id),
label,
int(count),
round(share, 2),
"#{:02X}{:02X}{:02X}".format(*color),
]
)
rows.sort(key=lambda row: float(row[3]), reverse=True)
return rows
def write_class_csv(path: Path, rows: Iterable[Iterable[object]]) -> None:
with path.open("w", newline="", encoding="utf-8") as handle:
writer = csv.writer(handle)
writer.writerow(["class_id", "class_name", "pixels", "share_percent", "color"])
writer.writerows(rows)
|