hewumars
5/7/2019 - 6:41 AM

COCO数据集可视化

from pycocotools.coco import COCO
import cv2
import pandas as pd
import os


def showNimages(imageidFile, annFile, imageFile, resultFile):
    """
    :param imageidFile: 要查看的图片imageid,存储一列在csv文件里 (目前设计的imageid需要为6位数,如果少于6位数,可以在前面加多个0)
    :param annFile:使用的标注文件
    :param imageFile:要读取的image所在文件夹
    :param resultFile:画了标注之后的image存储文件夹
    :return:
    """
    data = pd.read_csv(imageidFile)
    list = data.values.tolist()
    image_id = []  # 存储的是要提取图片id
    for i in range(len(list)):
        image_id.append(str(list[i][0]).zfill(7))
    print(image_id)
    print(len(image_id))
    coco = COCO(annFile)

    for i in range(len(image_id)):
        img = coco.loadImgs(int(i+1))[0]
        image = cv2.imread(imageFile + img['file_name'])
        annIds = coco.getAnnIds(imgIds=img['id'], iscrowd=None)
        anns = coco.loadAnns(annIds)
        for n in range(len(anns)):
            x, y, w, h = anns[n]['bbox']
            x, y, w, h = int(x), int(y), int(w), int(h)
            # print(x, y, w, h)
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255))
        cv2.imwrite(os.path.join(resultFile,img['file_name']), image)

    print("生成图片存在{}".format(resultFile))

if __name__ == "__main__":
    imageidFile = '../data/cocoMY/val.csv'
    annFile = '../data/cocoMY/annotations/val.json'
    imageFile = '/home/mars/hewu/OneStageDetect/20190501_datasets/JPEGImages/'
    resultFile = '/home/mars/hewu/OneStageDetect/resultCOCOVis'
    showNimages(imageidFile,annFile,imageFile,resultFile)