coco数据可视化

"""
coco_visual.py
可视化COCO数据集

Usage:
python coco_visual.py 标注文件路径 图片文件夹 [-s=可视化结果保存路径]
"""

import argparse
import cv2
import os
import shutil
import asyncio
import numpy as np
from skimage import draw
from pycocotools.coco import COCO
from typing import Union, Optional, List, Sequence

TEXT_MARGIN = 2
VIS_COLOR = [(244, 67, 54),
             (233, 30, 99),
             (156, 39, 176),
             (103, 58, 183),
             (63, 81, 181),
             (33, 150, 243),
             (3, 169, 244),
             (0, 188, 212),
             (0, 150, 136),
             (76, 175, 80),
             (139, 195, 74),
             (205, 220, 57),
             (255, 235, 59),
             (255, 193, 7),
             (255, 152, 0),
             (255, 87, 34),
             (121, 85, 72),
             (158, 158, 158),
             (96, 125, 139)]


def parse_args():
    parser = argparse.ArgumentParser(description='Analyze COCO Data')
    parser.add_argument('anno_path', type=str, help='coco annotation path')
    parser.add_argument('image_dir', type=str, help='coco image dir')
    parser.add_argument('--save-path', '-s', type=str, help='path to save visual result')
    args = parser.parse_args()
    return args


def _imread(image: Union[str, np.ndarray]):
    """
    read image
    Args:
        image (): image path or image nd-array

    Returns:

    """
    return image if isinstance(image, np.ndarray) else cv2.imread(image)


def _format_polygons(polygons: Union[List, np.ndarray]) -> List:
    """
    format polygons to segmentation format
    Args:
        polygons ():

    Returns:

    """
    if isinstance(polygons, np.ndarray):
        polygons = polygons.tolist()
    ps = []
    for polygon in polygons:
        if isinstance(polygon[0], Sequence):
            ps.append([tuple(p) for p in polygon])
        else:
            if len(polygon) % 2 == 0:
                ps.append([(polygon[i], polygon[i + 1]) for i in range(len(polygon)) if i % 2 == 0])
    return ps


def _get_text_size(text: str,
                   font_face: Optional[int] = cv2.FONT_HERSHEY_DUPLEX,
                   font_scale: Optional[float] = 0.6,
                   thickness: Optional[int] = 1):
    """
    get text font size
    Args:
        text ():
        font_face ():
        font_scale ():
        thickness ():

    Returns:

    """
    size = cv2.getTextSize(text, font_face, font_scale, thickness)
    text_width = size[0][0]
    text_height = size[0][1]
    return text_width, text_height, size[1]


def _get_text_color():
    """get text visual color
    """
    return (255, 255, 255)


def _get_color_tuple(color: Union[str, Sequence], alpha: Optional[float] = None):
    """
    get text color tuple give in color obj
    Args:
        color ():
        alpha ():

    Returns:

    """
    if alpha is not None:
        alpha = min(255, int(alpha)) if alpha > 1 else int(alpha * 255)
    if isinstance(color, (int, float)):
        if color < 1:
            color = int(255 * color)
        color = min(255, int(color))
        if alpha is not None:
            return (color, color, color, alpha)
        return (color, color, color)
    elif isinstance(color, Sequence):
        color = tuple(color)[:4]
        if alpha is None:
            return color
        if len(color) == 3:
            color = color + (alpha,)
        elif len(color) == 4:
            color = color[:3] + (alpha)
        else:
            raise ValueError("color value invalid: ", color)
        return color
    elif isinstance(color, str):
        from PIL import ImageColor
        return ImageColor.getrgb(color)
    else:
        import numpy as np
        if isinstance(color, np.ndarray) and color.ndim == 1:
            return _get_color_tuple(color.tolist(), alpha=alpha)
        raise ValueError("color value invalid: ", color)


def imdraw_bbox(image: Union[str, np.ndarray],
                xmin: float,
                ymin: float,
                xmax: float,
                ymax: float,
                color: Optional[Union[str, Sequence]] = (244, 67, 54),
                thickness: Optional[int] = 1,
                display_str: Optional[str] = "",
                text_color: Optional[Union[str, Sequence]] = None,
                use_normalized_coordinates: Optional[bool] = False):
    """
    draw bbox on image
    Args:
        image ():
        xmin ():
        ymin ():
        xmax ():
        ymax ():
        color ():
        thickness ():
        display_str ():
        text_color ():
        use_normalized_coordinates ():

    Returns:

    """
    assert xmin <= xmax, f"xmin({xmin}) shouldn't be langer than xmax({xmax})!"
    assert ymin <= ymax, f"ymin({ymin}) shouldn't be langer than ymax({ymax})!"

    image = _imread(image)
    image = image.copy()

    if xmin == xmax or ymin == ymax:
        return image

    im_height, im_width = image.shape[:2]
    if use_normalized_coordinates:
        (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                      ymin * im_height, ymax * im_height)
    else:
        (left, right, top, bottom) = (xmin, xmax, ymin, ymax)

    (left, right, top, bottom) = int(left), int(right), int(top), int(bottom)

    color = _get_color_tuple(color)

    cv2.rectangle(image, (left, top), (right, bottom), color, thickness=thickness)
    if display_str == "":
        return image

    text_width, text_height, line_height = _get_text_size(display_str)
    text_left = left + TEXT_MARGIN
    text_top = top - TEXT_MARGIN

    if text_top < TEXT_MARGIN:
        text_top = bottom + TEXT_MARGIN

    if text_top + text_height + TEXT_MARGIN > im_height:
        text_top = top + TEXT_MARGIN

    text_color = _get_color_tuple(text_color) if text_color is not None else _get_text_color()
    image = imdraw_text(image, display_str, text_left, text_top, text_color=text_color, bg_color=color)
    return image


def imdraw_polygons(image: Union[str, np.ndarray],
                    polygons: Union[List, np.ndarray],
                    color: Optional[Union[Sequence, str]] = 'red',
                    alpha: Optional[float] = 0.45):
    """
    draw polygon mask on image
    Args:
        image ():
        polygons ():
        color ():
        alpha ():

    Returns:

    """
    image = _imread(image)
    if not polygons:
        return image

    polygons = np.array(_format_polygons(polygons))
    polygons = np.squeeze(polygons, axis=0)
    X, Y = polygons[:, 0], polygons[:, 1]
    rr, cc = draw.polygon(Y, X)

    rgb = _get_color_tuple(color)
    draw.set_color(image, [rr, cc], color=rgb, alpha=alpha)

    return image


def imdraw_text(image: Union[str, np.ndarray],
                text: Optional[str] = "-",
                x: Optional[int] = 0,
                y: Optional[int] = 0,
                font_scale: Optional[float] = 0.6,
                thickness: Optional[int] = 1,
                text_color: Optional[Union[Sequence, str]] = (255, 255, 255),
                with_bg: Optional[bool] = True,
                bg_color: Optional[Union[Sequence, str]] = (244, 67, 54),
                bg_alpha: Optional[Union[float, int]] = 0.6):
    """
    draw text on image
    Args:
        image ():
        text ():
        x ():
        y ():
        font_scale ():
        thickness ():
        text_color ():
        with_bg ():
        bg_color ():
        bg_alpha ():

    Returns:

    """
    image = _imread(image)
    (height, width) = image.shape[:2]
    text_color = _get_color_tuple(text_color)
    if with_bg:
        bg_color = _get_color_tuple(bg_color, bg_alpha)
        text_width, text_height, line_height = _get_text_size(text, font_scale=font_scale, thickness=thickness)
        xmin = int(max(0, x - TEXT_MARGIN))
        ymin = int(max(0, y - TEXT_MARGIN - text_height))
        xmax = int(min(width, x + text_width + TEXT_MARGIN))
        ymax = int(min(height, y + TEXT_MARGIN))

        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), bg_color, thickness=-1)
    cv2.putText(image, text, (x, y), cv2.FONT_HERSHEY_DUPLEX, font_scale, text_color, thickness, cv2.LINE_AA)
    return image


def reset_dir(d: str):
    """
    remove dir and create new one
    Args:
        d ():

    Returns:

    """
    if os.path.isdir(d):
        shutil.rmtree(d)
    os.makedirs(d)


async def single_image_vis(img_info: dict,
                           image_dir: str,
                           annos: List,
                           cat_map: dict,
                           color_map: dict,
                           save_path: Optional[str] = None):
    image = _imread(os.path.join(image_dir, img_info["file_name"]))

    for ann in annos:
        cid = ann["category_id"]
        xmin, ymin = ann["bbox"][:2]
        xmax, ymax = xmin + ann["bbox"][2], ymin + ann["bbox"][3]
        segm = ann["segmentation"] if "segmentation" in ann else None
        image = imdraw_bbox(image, xmin, ymin, xmax, ymax, color=color_map[cid], 
                                display_str=cat_map[cid])
        image = imdraw_polygons(image, segm,  color=color_map[cid])

    if save_path:
        cv2.imwrite(save_path, image)


def visual_coco(anno_path: str, image_dir: str, save_path: str):
    """
    visual coco dataset
    Args:
        anno_path (): coco annotation file path
        image_dir (): coco image dir
        save_path (): visual result save dir

    Returns:

    """
    assert os.path.isfile(anno_path), f"coco annotation file({anno_path}) not exist"
    assert os.path.isdir(image_dir), f"coco image dir({image_dir}) not exist"

    coco = COCO(anno_path)
    # reset_dir(save_path)

    color_map = {c["id"]: VIS_COLOR[i % len(VIS_COLOR)] for i, c in enumerate(coco.cats.values())}
    cat_map = {id: cat["name"] for id, cat in coco.cats.items()}

    loop = asyncio.get_event_loop()
    tasks = []
    for id, img in coco.imgs.items():
        tasks.append(loop.create_task(single_image_vis(img,
                                                       image_dir,
                                                       coco.imgToAnns[img["id"]],
                                                       cat_map=cat_map,
                                                       color_map=color_map,
                                                       save_path=os.path.join(save_path, img["file_name"].split('/')[-1])
                                                       )))

    loop.run_until_complete(asyncio.wait(tasks))
    loop.close()

    print("\n visual coco dataset finish, saved at %s" % save_path)


def main():
    args = parse_args()
    visual_coco(
        anno_path=args.anno_path,
        image_dir=args.image_dir,
        save_path=args.save_path
    )


if __name__ == '__main__':
    main()