"""
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, SequenceTEXT_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 argsdef _imread(image: Union[str, np.ndarray]):"""read imageArgs:image (): image path or image nd-arrayReturns:"""return image if isinstance(image, np.ndarray) else cv2.imread(image)def _format_polygons(polygons: Union[List, np.ndarray]) -> List:"""format polygons to segmentation formatArgs: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 psdef _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 sizeArgs: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 objArgs: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 colorif len(color) == 3:color = color + (alpha,)elif len(color) == 4:color = color[:3] + (alpha)else:raise ValueError("color value invalid: ", color)return colorelif isinstance(color, str):from PIL import ImageColorreturn ImageColor.getrgb(color)else:import numpy as npif 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 imageArgs: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 imageim_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 imagetext_width, text_height, line_height = _get_text_size(display_str)text_left = left + TEXT_MARGINtext_top = top - TEXT_MARGINif text_top < TEXT_MARGIN:text_top = bottom + TEXT_MARGINif text_top + text_height + TEXT_MARGIN > im_height:text_top = top + TEXT_MARGINtext_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 imagedef 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 imageArgs:image ():polygons ():color ():alpha ():Returns:"""image = _imread(image)if not polygons:return imagepolygons = 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 imagedef 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 imageArgs: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 imagedef reset_dir(d: str):"""remove dir and create new oneArgs: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 Noneimage = 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 datasetArgs:anno_path (): coco annotation file pathimage_dir (): coco image dirsave_path (): visual result save dirReturns:"""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()
本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！
