Detectron2:Visualization
Detectron2에서 Predict 결과를 그리는 Python 코드. OpenCV의 Drawing API를 사용한다.
draw.py
심볼을 찾을 수 없습니다. 경고를 무시하기 위한 유틸리티. 파이썬의 Drawing 관련 함수는 이 패키지에서 호출한다.
# -*- coding: utf-8 -*-
from typing import Final, List, Tuple
from cv2 import FILLED as cv2_FILLED # noqa
from cv2 import FONT_HERSHEY_SIMPLEX as cv2_FONT_HERSHEY_SIMPLEX # noqa
from cv2 import LINE_4 as cv2_LINE_4 # noqa
from cv2 import LINE_8 as cv2_LINE_8 # noqa
from cv2 import LINE_AA as cv2_LINE_AA # noqa
from cv2 import addWeighted as cv2_addWeighted # noqa
from cv2 import fillPoly as cv2_fillPoly # noqa
from cv2 import getTextSize as cv2_getTextSize # noqa
from cv2 import polylines as cv2_polylines # noqa
from cv2 import putText as cv2_putText # noqa
from cv2 import rectangle as cv2_rectangle # noqa
from numpy import ndarray
LINE_4: Final[int] = cv2_LINE_4
LINE_8: Final[int] = cv2_LINE_8
LINE_AA: Final[int] = cv2_LINE_AA
FILLED: Final[int] = cv2_FILLED
FONT_HERSHEY_SIMPLEX: Final[int] = cv2_FONT_HERSHEY_SIMPLEX
DEFAULT_COLOR: Final[Tuple[int, int, int]] = (0, 0, 255)
DEFAULT_THICKNESS: Final[int] = 2
DEFAULT_LINE_TYPE: Final[int] = LINE_AA
DEFAULT_FONT: Final[int] = FONT_HERSHEY_SIMPLEX
DEFAULT_FONT_SCALE: Final[float] = 1.0
def fill_poly(
image: ndarray,
points: List[ndarray],
color=DEFAULT_COLOR,
line_type=DEFAULT_LINE_TYPE,
) -> None:
cv2_fillPoly(image, points, color, line_type)
def polylines(
image: ndarray,
points: List[ndarray],
closed=True,
color=DEFAULT_COLOR,
thickness=DEFAULT_THICKNESS,
line_type=DEFAULT_LINE_TYPE,
) -> None:
cv2_polylines(image, points, closed, color, thickness, line_type)
def add_weighted(
src1: ndarray,
alpha: float,
src2: ndarray,
beta: float,
gamma=0.0,
) -> ndarray:
return cv2_addWeighted(src1, alpha, src2, beta, gamma)
def rectangle(
image: ndarray,
point1: Tuple[int, int],
point2: Tuple[int, int],
color=DEFAULT_COLOR,
thickness=DEFAULT_THICKNESS,
line_type=DEFAULT_LINE_TYPE,
) -> None:
cv2_rectangle(image, point1, point2, color, thickness, line_type)
def put_text(
image: ndarray,
text: str,
bottom_left_corner: Tuple[int, int],
font_face=DEFAULT_FONT,
font_scale=DEFAULT_FONT_SCALE,
color=DEFAULT_COLOR,
thickness=DEFAULT_THICKNESS,
line_type=DEFAULT_LINE_TYPE,
) -> None:
cv2_putText(
image,
text,
bottom_left_corner,
font_face,
font_scale,
color,
thickness,
line_type,
bottomLeftOrigin=False,
)
def get_text_size(
text: str,
font_face=DEFAULT_FONT,
font_scale=DEFAULT_FONT_SCALE,
thickness=DEFAULT_THICKNESS,
) -> Tuple[Tuple[int, int], int]:
return cv2_getTextSize(text, font_face, font_scale, thickness)
adjust.py
색상을 어둡게 하거나 밝게 하기 위한 함수. Python:colorsys 패키지 사용.
# -*- coding: utf-8 -*-
from colorsys import hls_to_rgb, rgb_to_hls
from typing import Tuple
def adjust_color(r: int, g: int, b: int, factor: float) -> Tuple[int, int, int]:
h, l, s = rgb_to_hls(r / 255.0, g / 255.0, b / 255.0)
adjust_l = max(min(l * factor, 1.0), 0.0)
fr, fg, fb = hls_to_rgb(h, adjust_l, s)
return int(fr * 255), int(fg * 255), int(fb * 255)
def lighten_color(r: int, g: int, b: int, factor=0.1):
return adjust_color(r, g, b, 1 + factor)
def darken_color(r: int, g: int, b: int, factor=0.1):
return adjust_color(r, g, b, 1 - factor)
invert.py
색상을 반전시키기 위한 함수.
# -*- coding: utf-8 -*-
from typing import Tuple
def invert_color(r: int, g: int, b: int) -> Tuple[int, int, int]:
return abs(r - 255), abs(g - 255), abs(b - 255)
american_palette.py
미리 정의한 색상 팔래트 https://flatuicolors.com/palette/us 사이트 참조.
# -*- coding: utf-8 -*-
# https://flatuicolors.com/palette/us
# American Palette by Kevin Yang from
from typing import Tuple
def _rgb_to_bgr(r: int, g: int, b: int) -> Tuple[int, int, int]:
return b, g, r
LIGHT_GREENISH_BLUE = _rgb_to_bgr(85, 239, 196)
FADED_POSTER = _rgb_to_bgr(129, 236, 236)
GREEN_DARNER_TAIL = _rgb_to_bgr(116, 185, 255)
SHY_MOMENT = _rgb_to_bgr(162, 155, 254)
CITY_LIGHTS = _rgb_to_bgr(223, 230, 233)
MINT_LEAF = _rgb_to_bgr(0, 184, 148)
ROBINS_EGG_BLUE = _rgb_to_bgr(0, 206, 201)
ELECTRON_BLUE = _rgb_to_bgr(9, 132, 227)
EXODUS_FRUIT = _rgb_to_bgr(108, 92, 231)
SOOTHING_BREEZE = _rgb_to_bgr(178, 190, 195)
SOUR_LEMON = _rgb_to_bgr(255, 234, 167)
FIRST_DATE = _rgb_to_bgr(250, 177, 160)
PINK_GLAMOUR = _rgb_to_bgr(255, 118, 117)
PICO_8_PINK = _rgb_to_bgr(253, 121, 168)
AMERICAN_RIVER = _rgb_to_bgr(99, 110, 114)
BRIGHT_YARROW = _rgb_to_bgr(253, 203, 110)
ORANGEVILLE = _rgb_to_bgr(225, 112, 85)
CHI_GONG = _rgb_to_bgr(214, 48, 49)
PRUNUS_AVIUM = _rgb_to_bgr(232, 67, 147)
DRACULA_ORCHID = _rgb_to_bgr(45, 52, 54)
AMERICAN_PALETTE = [
LIGHT_GREENISH_BLUE,
FADED_POSTER,
GREEN_DARNER_TAIL,
SHY_MOMENT,
CITY_LIGHTS,
MINT_LEAF,
ROBINS_EGG_BLUE,
ELECTRON_BLUE,
EXODUS_FRUIT,
SOOTHING_BREEZE,
SOUR_LEMON,
FIRST_DATE,
PINK_GLAMOUR,
PICO_8_PINK,
AMERICAN_RIVER,
BRIGHT_YARROW,
ORANGEVILLE,
CHI_GONG,
PRUNUS_AVIUM,
DRACULA_ORCHID,
]
Predict
# -*- coding: utf-8 -*-
from dataclasses import dataclass
from io import StringIO
from typing import Final, List, Optional, Tuple
from adjust import darken_color
from american_palette import AMERICAN_PALETTE
from invert import invert_color
from draw import (
FILLED,
add_weighted,
fill_poly,
get_text_size,
polylines,
put_text,
rectangle,
)
from numpy import array, float64, int32, ndarray
Polyline = List[float] # x0, y0, x1, y1, x2, y2, ... xN, yN
Polylines = List[Polyline]
PolylinesList = List[Polylines]
Box = List[float] # x1, y1, x2, y2
Boxes = List[Box]
BLACK_COLOR: Final[Tuple[int, int, int]] = (0, 0, 0)
WHITE_COLOR: Final[Tuple[int, int, int]] = (255, 255, 255)
DEFAULT_LINE_WIDTH: Final[int] = 2
DEFAULT_OVERLAY_ALPHA: Final[float] = 0.3
DEFAULT_FONT_SCALE: Final[float] = 0.7
DEFAULT_LABEL_BOX_PADDING: Final[int] = 2
DEFAULT_DARKEN_FACTOR: Final[float] = 0.6
@dataclass
class Point:
x: float
y: float
@dataclass
class Rect:
x1: float
y1: float
x2: float
y2: float
def to_string(self):
return f"{self.__class__.__name__}[{self.x1},{self.y1},{self.x2},{self.y2}]"
def __repr__(self):
return self.to_string()
def __str__(self):
return self.to_string()
@property
def left(self):
return min(self.x1, self.x2)
@property
def right(self):
return max(self.x1, self.x2)
@property
def top(self):
return min(self.y1, self.y2)
@property
def bottom(self):
return max(self.y1, self.y2)
@property
def width(self):
return abs(self.x2 - self.x1)
@property
def height(self):
return abs(self.y2 - self.y1)
def intersection(self, obj: "Rect") -> bool:
x1 = max(self.left, obj.left)
y1 = max(self.top, obj.top)
x2 = min(self.right, obj.right)
y2 = min(self.bottom, obj.bottom)
return x1 < x2 and y1 < y2
@dataclass
class ColorInfo:
line_bgr: Tuple[int, int, int]
fill_bgr: Tuple[int, int, int]
label_bgr: Tuple[int, int, int]
@classmethod
def from_base_bgr(cls, base_color: Tuple[int, int, int], darken_factor: float):
fill_bgr = base_color
b, g, r = fill_bgr
line_bgr = darken_color(r, g, b, darken_factor)[::-1]
label_bgr = invert_color(line_bgr[0], line_bgr[1], line_bgr[2])
return cls(line_bgr, fill_bgr, label_bgr)
@dataclass
class LabelBoxInfo:
label_box_left: int = 0
label_box_top: int = 0
label_box_right: int = 0
label_box_bottom: int = 0
label_left: int = 0
label_bottom: int = 0
def as_label_box_left_top(self) -> Tuple[int, int]:
return self.label_box_left, self.label_box_top
def as_label_box_right_bottom(self) -> Tuple[int, int]:
return self.label_box_right, self.label_box_bottom
def as_label_left_bottom(self) -> Tuple[int, int]:
return self.label_left, self.label_bottom
def adjust_screen_position(
self,
width: Optional[int] = None,
height: Optional[int] = None,
) -> None:
"""If the 'label box' is off the screen, adjust its position."""
if self.label_box_left < 0:
# Move to right
move_right = abs(self.label_box_left)
self.label_box_left += move_right
self.label_box_right += move_right
self.label_left += move_right
elif width is not None and self.label_box_right > width:
# Move to left
move_left = self.label_box_right - width
assert move_left > 0
self.label_box_left -= move_left
self.label_box_right -= move_left
self.label_left -= move_left
if self.label_box_top < 0:
# Move to down
move_down = abs(self.label_box_top)
self.label_box_top += move_down
self.label_box_bottom += move_down
self.label_bottom += move_down
elif height is not None and self.label_box_bottom > height:
# Move to up
move_up = self.label_box_bottom - height
assert move_up > 0
self.label_box_top -= move_up
self.label_box_bottom -= move_up
self.label_bottom -= move_up
@dataclass
class PredictInstance:
index: int # aka class
score: float
label: Optional[str] = None
box: Optional[Rect] = None
polylines: Optional[Polylines] = None
def to_string(self):
io = StringIO()
io.write(self.__class__.__name__)
io.write(f"[index={self.index},score={self.score:.2f}")
if self.label is not None:
io.write(f",label={self.label}")
if self.box is not None:
io.write(f",box={self.box}")
io.write("]")
return io.getvalue()
def __repr__(self):
return self.to_string()
def __str__(self):
return self.to_string()
def to_drawable_label(self) -> str:
if self.label:
return "{} {:.0f}%".format(self.label, self.score * 100)
else:
return "#{} {:.0f}%".format(self.index, self.score * 100)
def as_numpy_polylines(self) -> List[ndarray]:
if not self.polylines:
return list()
if not isinstance(self.polylines[0], list):
element_typename = type(self.polylines[0]).__name__
raise TypeError(f"Unexpected polylines element: {element_typename}")
result = []
for polyline in self.polylines:
assert isinstance(polyline, list)
arr = array(polyline, float64).astype(int32).reshape((-1, 2))
result.append(arr)
return result
def get_text_size(self, font_scale: float) -> Tuple[int, int, int]:
label_size = get_text_size(self.to_drawable_label(), font_scale=font_scale)
label_width, label_height = label_size[0]
label_base_line = label_size[1]
return label_width, label_height, label_base_line
def get_label_box_info(
self,
font_scale: float,
label_box_padding: int,
image_width: Optional[int] = None,
image_height: Optional[int] = None,
) -> LabelBoxInfo:
label_width, label_height, label_base_line = self.get_text_size(font_scale)
box_left = int(self.box.x1) if self.box else 0
box_top = int(self.box.y1) if self.box else 0
label_box_left = box_left
label_box_top = box_top - label_height - (label_base_line * 2)
label_box_right = box_left + label_width + (label_box_padding * 2)
label_box_bottom = box_top
label_left = box_left + label_box_padding
label_bottom = box_top - label_base_line
result = LabelBoxInfo(
label_box_left,
label_box_top,
label_box_right,
label_box_bottom,
label_left,
label_bottom,
)
result.adjust_screen_position(image_width, image_height)
return result
def draw_polylines(
self,
image: ndarray,
line_width: int,
fill_color: Tuple[int, int, int],
line_color: Tuple[int, int, int],
) -> None:
numpy_polylines = self.as_numpy_polylines()
fill_poly(image, numpy_polylines, fill_color)
polylines(image, numpy_polylines, True, line_color, line_width)
def draw_box(
self,
image: ndarray,
line_width: int,
line_color: Tuple[int, int, int],
) -> None:
box_left = int(self.box.x1) if self.box else 0
box_right = int(self.box.x2) if self.box else 0
box_top = int(self.box.y1) if self.box else 0
box_bottom = int(self.box.y2) if self.box else 0
rectangle(
image,
(box_left, box_top),
(box_right, box_bottom),
line_color,
line_width,
)
def draw_label_box(
self,
image: ndarray,
font_scale: float,
label_box_padding: int,
fill_color: Tuple[int, int, int],
) -> None:
label_box_info = self.get_label_box_info(
font_scale=font_scale,
label_box_padding=label_box_padding,
image_width=image.shape[1],
image_height=image.shape[0],
)
rectangle(
image,
label_box_info.as_label_box_left_top(),
label_box_info.as_label_box_right_bottom(),
fill_color,
FILLED,
)
def draw_label(
self,
image: ndarray,
font_scale: float,
label_box_padding: int,
label_color: Tuple[int, int, int],
) -> None:
label_box_info = self.get_label_box_info(
font_scale=font_scale,
label_box_padding=label_box_padding,
image_width=image.shape[1],
image_height=image.shape[0],
)
put_text(
image,
self.to_drawable_label(),
label_box_info.as_label_left_bottom(),
font_scale=font_scale,
color=label_color,
)
def draw_overlay(
self,
image: ndarray,
fill_color: Tuple[int, int, int],
line_color: Tuple[int, int, int],
line_width: int,
overlay_alpha: float,
font_scale: float,
label_box_padding: int,
with_polylines=True,
with_box=True,
with_label=True,
) -> None:
overlay = image.copy()
if with_polylines and self.polylines:
self.draw_polylines(overlay, line_width, fill_color, line_color)
if with_box and self.box is not None:
self.draw_box(overlay, line_width, line_color)
if with_label and self.label is not None and self.box is not None:
self.draw_label_box(overlay, font_scale, label_box_padding, line_color)
alpha = 1 - overlay_alpha
image[:] = add_weighted(overlay, alpha, image, overlay_alpha)
def draw(
self,
image: ndarray,
line_width=DEFAULT_LINE_WIDTH,
overlay_alpha=DEFAULT_OVERLAY_ALPHA,
font_scale=DEFAULT_FONT_SCALE,
label_box_padding=DEFAULT_LABEL_BOX_PADDING,
darken_factor=DEFAULT_DARKEN_FACTOR,
with_polylines=True,
with_box=True,
with_label=True,
palette: List[Tuple[int, int, int]] = AMERICAN_PALETTE,
) -> None:
color_info = ColorInfo.from_base_bgr(palette[self.index], darken_factor)
self.draw_overlay(
image=image,
fill_color=color_info.fill_bgr,
line_color=color_info.line_bgr,
line_width=line_width,
overlay_alpha=overlay_alpha,
font_scale=font_scale,
label_box_padding=label_box_padding,
with_polylines=with_polylines,
with_box=with_box,
with_label=with_label,
)
# If the label is printed as an overlay, readability is reduced,
# so it should be drawn opaque.
self.draw_label(
image=image,
font_scale=font_scale,
label_box_padding=label_box_padding,
label_color=color_info.label_bgr,
)
@dataclass
class PredictQ:
threshold: Optional[float] = None
roi: Optional[Rect] = None
@dataclass
class PredictA:
predictions: List[PredictInstance]
iteration: Optional[int] = None
elapsed: Optional[float] = None
def to_string(self):
io = StringIO()
io.write(self.__class__.__name__)
io.write(f"[predictions={len(self.predictions)}")
if self.iteration is not None:
io.write(f",iteration={self.iteration}")
if self.elapsed is not None:
io.write(f",elapsed={self.elapsed:.2f}s")
io.write("]")
return io.getvalue()
def __repr__(self):
return self.to_string()
def __str__(self):
return self.to_string()
@classmethod
def from_raws(
cls,
classes: List[int],
labels: List[str],
scores: List[float],
boxes: Boxes,
polylines_list: PolylinesList,
iteration: Optional[int] = None,
elapsed: Optional[float] = None,
):
len_labels = len(labels)
len_scores = len(scores)
len_boxes = len(boxes)
len_polylines_list = len(polylines_list)
if not (len_labels == len_scores == len_boxes == len_polylines_list):
raise IndexError("All raw data must be the same length")
predictions = list()
for i in range(len_labels):
x1, y1, x2, y2 = boxes[i]
instance = PredictInstance(
index=classes[i],
score=scores[i],
label=labels[i],
box=Rect(x1, y1, x2, y2),
polylines=polylines_list[i],
)
predictions.append(instance)
return cls(
predictions=predictions,
iteration=iteration,
elapsed=elapsed,
)
def draw(
self,
image: ndarray,
line_width=DEFAULT_LINE_WIDTH,
overlay_alpha=DEFAULT_OVERLAY_ALPHA,
font_scale=DEFAULT_FONT_SCALE,
label_box_padding=DEFAULT_LABEL_BOX_PADDING,
darken_factor=DEFAULT_DARKEN_FACTOR,
with_polylines=True,
with_box=True,
with_label=True,
palette: List[Tuple[int, int, int]] = AMERICAN_PALETTE,
) -> None:
# If the printed 'label' overlaps,
# an object with a higher score is drawn on top.
predictions = self.predictions.copy()
predictions.sort(key=lambda x: x.score)
for predict in predictions:
predict.draw(
image=image,
line_width=line_width,
overlay_alpha=overlay_alpha,
font_scale=font_scale,
label_box_padding=label_box_padding,
darken_factor=darken_factor,
with_polylines=with_polylines,
with_box=with_box,
with_label=with_label,
palette=palette,
)
Usage
다음과 같이 사용한다:
from io import StringIO
from typing import List, Optional
def create_text_labels(
classes: List[int],
scores: Optional[List[float]] = None,
is_crowd: Optional[List[bool]] = None,
*,
labels: Optional[List[str]] = None,
) -> List[str]:
result = list()
for i, class_number in enumerate(classes):
io = StringIO()
io.write(labels[class_number] if labels else f"[{class_number}]")
if scores:
io.write(" {:.0f}%".format(scores[i] * 100))
if is_crowd and is_crowd[i]:
io.write("|crowd")
result.append(io.getvalue())
return result
def read_labels_file(path: str) -> List[str]:
try:
with open(path, mode="r") as f:
return f.read().strip().split("\n")
except BaseException: # noqa
return list()
cpu_device = device("cpu")
label_table = read_labels_file(labels_file if labels_file else str())
def predictions_to_answer(
predictions: Dict[str, Any],
width: int,
height: int,
) -> PredictA:
assert "instances" in predictions
objs = predictions["instances"].to(cpu_device)
assert isinstance(objs, Instances)
assert objs.has("pred_boxes")
assert isinstance(objs.pred_boxes, Boxes)
boxes = [b.tolist() for b in objs.pred_boxes]
assert isinstance(boxes, list)
assert objs.has("scores")
assert isinstance(objs.scores, Tensor)
scores = objs.scores.tolist()
assert isinstance(scores, list)
assert objs.has("pred_classes")
classes = objs.pred_classes.tolist()
assert isinstance(classes, list)
labels = create_text_labels(
classes=classes,
scores=None,
is_crowd=None,
labels=label_table,
)
# objs.has("pred_keypoints")
# keypoints = objs.pred_keypoints if objs.has("pred_keypoints") else None
# assert keypoints is None
assert objs.has("pred_masks")
pred_masks = asarray(objs.pred_masks)
masks = [GenericMask(x, height, width) for x in pred_masks]
polylines_list = list()
for mask in masks:
polylines = list()
for polygon in mask.polygons:
raw_polyline = polygon.tolist()
assert isinstance(raw_polyline, list)
polylines.append(raw_polyline)
polylines_list.append(polylines)
return PredictA.from_raws(
classes,
labels,
scores,
boxes,
polylines_list,
)
cfg = ...
image = ...
height, width, channel = image.shape
assert height >= 1
assert width >= 1
assert channel == 3
from detectron2.engine.defaults import DefaultPredictor
predictor = DefaultPredictor(cfg)
predictions = predictor(image)
answer = predictions_to_answer(predictions, width, height)
answer.draw(image)