【Code-coco检测框显示】

import json
import os
import cv2
import numpy as np
import random
import math
import matplotlib.pyplot as plt
from pycocotools.coco import COCO
%matplotlib inline
def visualize(image_dir, annotation_file, file_name,coco_dt):
    '''
    Args:
        image_dir (str): image directory
        annotation_file (str): annotation (.json) file path
        file_name (str): target file name (.jpg)
    Returns:
        None
    Example:
        image_dir = "./images"
        annotation_file = "./annotations.json"
        file_name = 'img_0028580.jpg'
        visualize(image_dir, annotation_file, file_name)
    '''
    image_path = os.path.join( image_dir, file_name )
    assert os.path.exists( image_path ), "image path not exist."
    assert os.path.exists( annotation_file ), "annotation file path not exist"
    image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
    with open(annotation_file) as f:
        data = json.load(f)
    image_id = None
    for i in data['images']:
        if i['file_name'] == file_name:
            image_id = i['id']
            break
    if not image_id:
        print("file name {} not found.".format(file_name))
    large_img = True if max( image.shape[0], image.shape[1] ) > 1000 else False
    linewidth = 10 if large_img else 10
    annots = []
    for a in data['annotations']:
        if a['image_id'] == image_id:
            bbox = [int(b) for b in a['bbox']]
            annots.append(bbox)

    sample = {'img': image, 'annot': np.array(annots)}
# #     sample = RandomFlip()(sample) # 这里可以做数据增强
    image = sample['img']
    annots = sample['annot']
    for bbox in annots:
        
        bbox[2] = bbox[2] + bbox[0] - 1
        bbox[3] = bbox[3] + bbox[1] - 1
        print("GT:",(bbox[0], bbox[1]), (bbox[2], bbox[3]))
        cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), linewidth)
    ####
    annots_dt = coco_dt.loadAnns(coco_dt.getAnnIds(image_id))
    for ann in annots_dt:
        if ann['score'] < 0.2:
            continue
        bbox = ann['bbox']
        score = ann['score']
        bbox[2] = bbox[2] + bbox[0] - 1
        bbox[3] = bbox[3] + bbox[1] - 1
        obj = ann['category_id']
        cv2.putText(image, '{}, {:.3f}'.format(obj, score),
                (int(bbox[0]), int(bbox[1]) + 10), cv2.FONT_HERSHEY_SIMPLEX,4,
                (255, 255, 0),4)
        print("DT:",(bbox[0], bbox[1]), (bbox[2], bbox[3]),score)
        
        cv2.rectangle(image, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0,0,255), linewidth)
    ###
    if large_img:
        plt.figure(figsize=(8,6))
    else:
        plt.figure(figsize=(5,5))
    plt.imshow(image)
    plt.show()
    return

images_root_path = 'datasets/images/'
image_list = os.listdir(images_root_path)
ann_json = 'instances_val2017.json'
visualize(images_root_path,ann_json,image_list[67])
dt_json = "predictBox.json"
coco_gt = COCO(ann_json)
coco_dt = coco_gt.loadRes(dt_json)
visualize(images_root_path,ann_json,image_list[1],coco_dt)

数据增强样例

class Normalizer(object):

    def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
        self.mean = np.array([[mean]])
        self.std = np.array([[std]])

    def __call__(self, sample):
        image, annots = sample['img'], sample['annot']

        return {'img': ((image.astype(np.float32) - self.mean) / self.std), 'annot': annots}