LaDeco / core.py
lichih's picture
Go
cdd6c01 verified
"""Core part of LaDeco v2
Example usage:
>>> from core import Ladeco
>>> from PIL import Image
>>> from pathlib import Path
>>>
>>> # predict
>>> ldc = Ladeco()
>>> imgs = (thing for thing in Path("example").glob("*.jpg"))
>>> out = ldc.predict(imgs)
>>>
>>> # output - visualization
>>> segs = out.visualize(level=2)
>>> segs[0].image.show()
>>>
>>> # output - element area
>>> area = out.area()
>>> area[0]
{"fid": "example/.jpg", "l1_nature": 0.673, "l1_man_made": 0.241, ...}
"""
from matplotlib.patches import Rectangle
from pathlib import Path
from PIL import Image
from transformers import AutoModelForUniversalSegmentation, AutoProcessor
import math
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import torch
from functools import lru_cache
from matplotlib.figure import Figure
import numpy.typing as npt
from typing import Iterable, NamedTuple, Generator
from tqdm import tqdm
class LadecoVisualization(NamedTuple):
filename: str
image: Figure
class Ladeco:
def __init__(self,
model_name: str = "shi-labs/oneformer_ade20k_swin_large",
area_threshold: float = 0.01,
device: str | None = None,
):
if device is None:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
else:
self.device = device
self.processor = AutoProcessor.from_pretrained(model_name)
self.model = AutoModelForUniversalSegmentation.from_pretrained(model_name).to(self.device)
self.area_threshold = area_threshold
self.ade20k_labels = {
name.strip(): int(idx)
for name, idx in self.model.config.label2id.items()
}
self.ladeco2ade20k: dict[str, tuple[int]] = _get_ladeco_labels(self.ade20k_labels)
def predict(
self, image_paths: str | Path | Iterable[str | Path], show_progress: bool = False
) -> "LadecoOutput":
if isinstance(image_paths, (str, Path)):
imgpaths = [image_paths]
else:
imgpaths = list(image_paths)
images = (
Image.open(img_path).convert("RGB")
for img_path in imgpaths
)
# batch inference functionality of OneFormer is broken
masks: list[torch.Tensor] = []
for img in tqdm(images, total=len(imgpaths), desc="Segmenting", disable=not show_progress):
samples = self.processor(
images=img, task_inputs=["semantic"], return_tensors="pt"
).to(self.device)
with torch.no_grad():
outputs = self.model(**samples)
masks.append(
self.processor.post_process_semantic_segmentation(outputs)[0]
)
return LadecoOutput(imgpaths, masks, self.ladeco2ade20k, self.area_threshold)
class LadecoOutput:
def __init__(
self,
filenames: list[str | Path],
masks: torch.Tensor,
ladeco2ade: dict[str, tuple[int]],
threshold: float,
):
self.filenames = filenames
self.masks = masks
self.ladeco2ade: dict[str, tuple[int]] = ladeco2ade
self.ade2ladeco: dict[int, str] = {
idx: label
for label, indices in self.ladeco2ade.items()
for idx in indices
}
self.threshold = threshold
def visualize(self, level: int) -> list[LadecoVisualization]:
return list(self.ivisualize(level))
def ivisualize(self, level: int) -> Generator[LadecoVisualization, None, None]:
colormaps = self.color_map(level)
labelnames = [name for name in self.ladeco2ade if name.startswith(f"l{level}")]
for fname, mask in zip(self.filenames, self.masks):
size = mask.shape + (3,) # (H, W, RGB)
vis = torch.zeros(size, dtype=torch.uint8)
for name in labelnames:
for idx in self.ladeco2ade[name]:
color = torch.tensor(colormaps[name] * 255, dtype=torch.uint8)
vis[mask == idx] = color
with Image.open(fname) as img:
target_size = img.size
vis = Image.fromarray(vis.numpy(), mode="RGB").resize(target_size)
fig, ax = plt.subplots()
ax.imshow(vis)
ax.axis('off')
yield LadecoVisualization(filename=str(fname), image=fig)
def area(self) -> list[dict[str, float | str]]:
return list(self.iarea())
def iarea(self) -> Generator[dict[str, float | str], None, None]:
n_label_ADE20k = 150
for filename, mask in zip(self.filenames, self.masks):
ade_ratios = torch.tensor([(mask == i).count_nonzero() / mask.numel() for i in range(n_label_ADE20k)])
#breakpoint()
ldc_ratios: dict[str, float] = {
label: round(ade_ratios[list(ade_indices)].sum().item(), 4)
for label, ade_indices in self.ladeco2ade.items()
}
ldc_ratios: dict[str, float] = {
label: 0 if ratio < self.threshold else ratio
for label, ratio in ldc_ratios.items()
}
others = round(1 - ldc_ratios["l1_nature"] - ldc_ratios["l1_man_made"], 4)
nfi = round(ldc_ratios["l1_nature"]/ (ldc_ratios["l1_nature"] + ldc_ratios.get("l1_man_made", 0) + 1e-6), 4)
yield {
"fid": str(filename), **ldc_ratios, "others": others, "LC_NFI": nfi,
}
def color_map(self, level: int) -> dict[str, npt.NDArray[np.float64]]:
"returns {'label_name': (R, G, B), ...}, where (R, G, B) in range [0, 1]"
labels = [
name for name in self.ladeco2ade.keys() if name.startswith(f"l{level}")
]
if len(labels) == 0:
raise RuntimeError(
f"LaDeco only has 4 levels in 1, 2, 3, 4. You assigned {level}."
)
colormap = mpl.colormaps["viridis"].resampled(len(labels)).colors[:, :-1]
# [:, :-1]: discard alpha channel
return {name: color for name, color in zip(labels, colormap)}
def color_legend(self, level: int) -> Figure:
colors = self.color_map(level)
match level:
case 1:
ncols = 1
case 2:
ncols = 1
case 3:
ncols = 2
case 4:
ncols = 5
cell_width = 212
cell_height = 22
swatch_width = 48
margin = 12
nrows = math.ceil(len(colors) / ncols)
width = cell_width * ncols + 2 * margin
height = cell_height * nrows + 2 * margin
dpi = 72
fig, ax = plt.subplots(figsize=(width / dpi, height / dpi), dpi=dpi)
fig.subplots_adjust(margin/width, margin/height,
(width-margin)/width, (height-margin*2)/height)
ax.set_xlim(0, cell_width * ncols)
ax.set_ylim(cell_height * (nrows-0.5), -cell_height/2.)
ax.yaxis.set_visible(False)
ax.xaxis.set_visible(False)
ax.set_axis_off()
for i, name in enumerate(colors):
row = i % nrows
col = i // nrows
y = row * cell_height
swatch_start_x = cell_width * col
text_pos_x = cell_width * col + swatch_width + 7
ax.text(text_pos_x, y, name, fontsize=14,
horizontalalignment='left',
verticalalignment='center')
ax.add_patch(
Rectangle(xy=(swatch_start_x, y-9), width=swatch_width,
height=18, facecolor=colors[name], edgecolor='0.7')
)
ax.set_title(f"LaDeco Color Legend - Level {level}")
return fig
def _get_ladeco_labels(ade20k: dict[str, int]) -> dict[str, tuple[int]]:
labels = {
# level 4 labels
# under l3_architecture
"l4_hovel": (ade20k["hovel, hut, hutch, shack, shanty"],),
"l4_building": (ade20k["building"], ade20k["house"]),
"l4_skyscraper": (ade20k["skyscraper"],),
"l4_tower": (ade20k["tower"],),
# under l3_archi_parts
"l4_step": (ade20k["step, stair"],),
"l4_canopy": (ade20k["awning, sunshade, sunblind"], ade20k["canopy"]),
"l4_arcade": (ade20k["arcade machine"],),
"l4_door": (ade20k["door"],),
"l4_window": (ade20k["window"],),
"l4_wall": (ade20k["wall"],),
# under l3_roadway
"l4_stairway": (ade20k["stairway, staircase"],),
"l4_sidewalk": (ade20k["sidewalk, pavement"],),
"l4_road": (ade20k["road, route"],),
# under l3_furniture
"l4_sculpture": (ade20k["sculpture"],),
"l4_flag": (ade20k["flag"],),
"l4_can": (ade20k["trash can"],),
"l4_chair": (ade20k["chair"],),
"l4_pot": (ade20k["pot"],),
"l4_booth": (ade20k["booth"],),
"l4_streetlight": (ade20k["street lamp"],),
"l4_bench": (ade20k["bench"],),
"l4_fence": (ade20k["fence"],),
"l4_table": (ade20k["table"],),
# under l3_vehicle
"l4_bike": (ade20k["bicycle"],),
"l4_motorbike": (ade20k["minibike, motorbike"],),
"l4_van": (ade20k["van"],),
"l4_truck": (ade20k["truck"],),
"l4_bus": (ade20k["bus"],),
"l4_car": (ade20k["car"],),
# under l3_sign
"l4_traffic_sign": (ade20k["traffic light"],),
"l4_poster": (ade20k["poster, posting, placard, notice, bill, card"],),
"l4_signboard": (ade20k["signboard, sign"],),
# under l3_vert_land
"l4_rock": (ade20k["rock, stone"],),
"l4_hill": (ade20k["hill"],),
"l4_mountain": (ade20k["mountain, mount"],),
# under l3_hori_land
"l4_ground": (ade20k["earth, ground"], ade20k["land, ground, soil"]),
"l4_field": (ade20k["field"],),
"l4_sand": (ade20k["sand"],),
"l4_dirt": (ade20k["dirt track"],),
"l4_path": (ade20k["path"],),
# under l3_flower
"l4_flower": (ade20k["flower"],),
# under l3_grass
"l4_grass": (ade20k["grass"],),
# under l3_shrub
"l4_flora": (ade20k["plant"],),
# under l3_arbor
"l4_tree": (ade20k["tree"],),
"l4_palm": (ade20k["palm, palm tree"],),
# under l3_hori_water
"l4_lake": (ade20k["lake"],),
"l4_pool": (ade20k["pool"],),
"l4_river": (ade20k["river"],),
"l4_sea": (ade20k["sea"],),
"l4_water": (ade20k["water"],),
# under l3_vert_water
"l4_fountain": (ade20k["fountain"],),
"l4_waterfall": (ade20k["falls"],),
# under l3_human
"l4_person": (ade20k["person"],),
# under l3_animal
"l4_animal": (ade20k["animal"],),
# under l3_sky
"l4_sky": (ade20k["sky"],),
}
labels = labels | {
# level 3 labels
# under l2_landform
"l3_hori_land": labels["l4_ground"] + labels["l4_field"] + labels["l4_sand"] + labels["l4_dirt"] + labels["l4_path"],
"l3_vert_land": labels["l4_mountain"] + labels["l4_hill"] + labels["l4_rock"],
# under l2_vegetation
"l3_woody_plant": labels["l4_tree"] + labels["l4_palm"] + labels["l4_flora"],
"l3_herb_plant": labels["l4_grass"],
"l3_flower": labels["l4_flower"],
# under l2_water
"l3_hori_water": labels["l4_water"] + labels["l4_sea"] + labels["l4_river"] + labels["l4_pool"] + labels["l4_lake"],
"l3_vert_water": labels["l4_fountain"] + labels["l4_waterfall"],
# under l2_bio
"l3_human": labels["l4_person"],
"l3_animal": labels["l4_animal"],
# under l2_sky
"l3_sky": labels["l4_sky"],
# under l2_archi
"l3_architecture": labels["l4_building"] + labels["l4_hovel"] + labels["l4_tower"] + labels["l4_skyscraper"],
"l3_archi_parts": labels["l4_wall"] + labels["l4_window"] + labels["l4_door"] + labels["l4_arcade"] + labels["l4_canopy"] + labels["l4_step"],
# under l2_street
"l3_roadway": labels["l4_road"] + labels["l4_sidewalk"] + labels["l4_stairway"],
"l3_furniture": labels["l4_table"] + labels["l4_chair"] + labels["l4_fence"] + labels["l4_bench"] + labels["l4_streetlight"] + labels["l4_booth"] + labels["l4_pot"] + labels["l4_can"] + labels["l4_flag"] + labels["l4_sculpture"],
"l3_vehicle": labels["l4_car"] + labels["l4_bus"] + labels["l4_truck"] + labels["l4_van"] + labels["l4_motorbike"] + labels["l4_bike"],
"l3_sign": labels["l4_signboard"] + labels["l4_poster"] + labels["l4_traffic_sign"],
}
labels = labels | {
# level 2 labels
# under l1_nature
"l2_landform": labels["l3_hori_land"] + labels["l3_vert_land"],
"l2_vegetation": labels["l3_woody_plant"] + labels["l3_herb_plant"] + labels["l3_flower"],
"l2_water": labels["l3_hori_water"] + labels["l3_vert_water"],
"l2_bio": labels["l3_human"] + labels["l3_animal"],
"l2_sky": labels["l3_sky"],
# under l1_man_made
"l2_archi": labels["l3_architecture"] + labels["l3_archi_parts"],
"l2_street": labels["l3_roadway"] + labels["l3_furniture"] + labels["l3_vehicle"] + labels["l3_sign"],
}
labels = labels | {
# level 1 labels
"l1_nature": labels["l2_landform"] + labels["l2_vegetation"] + labels["l2_water"] + labels["l2_bio"] + labels["l2_sky"],
"l1_man_made": labels["l2_archi"] + labels["l2_street"],
}
return labels
if __name__ == "__main__":
ldc = Ladeco()
image = Path("images") / "canyon_3011_00002354.jpg"
out = ldc.predict(image)