【画框脚本】YOLO和COCO格式画框

1. yolo格式画框

import cv2
import os
import colorsys
import random
from tqdm import tqdm

def get_n_hls_colors(num):
  hls_colors = []
  i = 0
  step = 360.0 / num
  while i < 360:
    h = i
    s = 90 + random.random() * 10
    l = 50 + random.random() * 10
    _hlsc = [h / 360.0, l / 100.0, s / 100.0]
    hls_colors.append(_hlsc)
    i += step

  return hls_colors

def ncolors(num):
  rgb_colors = []
  if num < 1:
    return rgb_colors
  hls_colors = get_n_hls_colors(num)
  for hlsc in hls_colors:
    _r, _g, _b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])
    r, g, b = [int(x * 255.0) for x in (_r, _g, _b)]
    rgb_colors.append([r, g, b])

  return rgb_colors

def convert(bbox,shape):
    x1 = int((bbox[0] - bbox[2] / 2.0) * shape[1])
    y1 = int((bbox[1] - bbox[3] / 2.0) * shape[0])
    x2 = int((bbox[0] + bbox[2] / 2.0) * shape[1])
    y2 = int((bbox[1] + bbox[3] / 2.0) * shape[0])
    return (x1,y1,x2,y2)

n = 80 # 类别数
# 获取n种区分度较大的rgb值
colors = ncolors(n)

images_list = os.listdir('images/') # 获取图片名列表
images_dir = 'images/' # 图片目录
labels_dir = 'labels/' # label目录
output_dir = 'showbbox/' # 输出图片目录

# import pdb;pdb.set_trace()
for img_id in tqdm(images_list):
    img = cv2.imread(images_dir + img_id) 
    # 判断后缀是为了排除隐藏文件.ipynb_checkpoints
    if img_id[-4:] != 'jpeg' and img_id[-3:] != 'jpg':
      continue
    shape = img.shape[0:2]
    txt_id = img_id.replace('jpeg', 'txt').replace('jpg', 'txt') 
    with open(labels_dir + txt_id) as r:
        lines = r.readlines()
        for line in lines:
            line = [float(i) for i in line.split(' ')] # 按空格划分并转换float类型
            label = int(line[0]) #获取类别信息
            bbox = line[1:] # 获取box信息
            (x1,y1,x2,y2) = convert(bbox,shape)
            # import pdb;pdb.set_trace()
            cv2.rectangle(img, (x1, y1), (x2, y2), (colors[label][2], colors[label][1], colors[label][0]), 3)
            # cv2.waitKey(0)
            cv2.putText(img, "{}".format(label),
                    (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    colors[label], 2)

    # print(output_dir + img_id)
    cv2.imwrite(output_dir + img_id, img)

2. COCO格式画框

import cv2
import argparse
import json
import os
from tqdm import tqdm
 
# ground-truth xywh
def gt_select_xywh(args):
    json_file = open(args.gt_json_path)
    infos = json.load(json_file)
    # import pdb;pdb.set_trace()
    images = infos["images"]
    annos = infos["annotations"]
    assert len(images) == len(images)
    for i in tqdm(range(len(images))):
        im_id = images[i]["id"]
        im_path = os.path.join(args.image_path, images[i]["file_name"])
        img = cv2.imread(im_path)
        for j in range(len(annos)):
            category_id = annos[j]["category_id"]
            if annos[j]["image_id"] == im_id:      
                x, y, w, h = annos[j]["bbox"]
                x, y, w, h = int(x), int(y), int(w), int(h)
                x2, y2 = x + w, y + h
                img = cv2.rectangle(img, (x, y), (x2, y2), (0, 255, 255), thickness=2)  # yellow
                img = cv2.putText(img, "{}".format(category_id),(x + 5, y + 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,127,0), 2)
                img_name = os.path.join(args.outpath, images[i]["file_name"].split('/')[-1])
                # import pdb;pdb.set_trace()
                cv2.imwrite(img_name, img)
                # continue
        # print(i)

# ground-truth xyxy
def gt_select_xyxy(args):
    json_file = open(args.gt_json_path)
    infos = json.load(json_file)
    images = infos["images"]
    annos = infos["annotations"]
    assert len(images) == len(images)
    for i in tqdm(range(len(images))):
        im_id = images[i]["id"]
        im_path = os.path.join(args.image_path, images[i]["file_name"])
        img = cv2.imread(im_path)
        for j in range(len(annos)):
            category_id = annos[j]["category_id"]
            if annos[j]["image_id"] == im_id:
                x1, y1, x2, y2 = annos[j]["bbox"]
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                # object_name = annos[j][""]
                img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 255), thickness=2)  # yellow
                img = cv2.putText(img, "{}".format(category_id),(x1 + 5, y1 + 5), cv2.FONT_HERSHEY_SIMPLEX, 2, (255,127,0), 2)
                img_name = os.path.join(args.outpath, images[i]["file_name"].split('/')[-1])
                # import pdb;pdb.set_trace()
                cv2.imwrite(img_name, img)
                # continue
        # print(i)

category_name = {0: 'person', 1:'child'}
# predict
def predict_select(args):
    infos = json.load(open(args.pred_json_path))
    pred_gt_file = json.load(open(args.gt_json_path))
    image_name_id_map = {i['id']: i['file_name'] for i in pred_gt_file['images']}
    # import pdb;pdb.set_trace()
    for i in tqdm(infos):
        im_path = os.path.join(args.image_path, image_name_id_map[i["image_id"]])
        img_name = os.path.join(args.outpath, image_name_id_map[i["image_id"]])
        score = str(i["score"])
        category = i['category_id']
        if not os.path.exists(img_name):
            img = cv2.imread(im_path)
        else: 
            img = cv2.imread(img_name)
        x, y, w, h = i["bbox"]
        x, y, w, h = int(x), int(y), int(w), int(h)
        x2, y2 = x + w, y + h
        if float(score) >= 0.25:
            cv2.rectangle(img, (x, y), (x2, y2), (0, 0, 255), thickness=2) # red
            cv2.putText(img, "{} {}".format(score, category),(x2, y2 + 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,127,0), 2)
        # else:
        #     cv2.rectangle(img, (x, y), (x2, y2), (255, 255, 0), 2)  # green
        img_name = os.path.join(args.outpath, image_name_id_map[i["image_id"]])
        cv2.imwrite(img_name, img)
    print("Done!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Start convert.')
    parser.add_argument('--pred_json_path', type=str, help='predition json file path')
    parser.add_argument('--gt_json_path', type=str, help='predition json file path')
    parser.add_argument('--image_path', type=str, help='raw image data dir')
    parser.add_argument('--outpath', type=str, help='image box dir')
    args = parser.parse_args()
    # predict_select(args) # predict json draw box
    gt_select_xywh(args) # gt json draw box (xywh)
    # gt_select_xyxy(args) # gt json draw box (xyxy)