目标检测一些函数

1. 计算bbox的iou

def myiou(bbox_A, bbox_B, mode='xyxy'): #xyxy
    x0_a,y0_a,x1_a,y1_a = bbox_A
    x0_b,y0_b,x1_b,y1_b = bbox_B
    areaA = (y1_a-y0_a) * (x1_a-x0_a)
    areaB = (y1_b-y0_b) * (x1_b-x0_b)
    overlap = max(min(x1_a,x1_b)-max(x0_a,x0_b), 0) * (min(y1_a,y1_b)-max(y0_a,y0_b))

    iou = overlap / (areaA+areaB-overlap+1e-6)
    return iou

2. 对文件夹中的图片按照整数序排序（非字典序）

def sort_files(imgs):
    sorted_imgs = []
    videos = []
    for img_name in imgs:
        videos.append(int(img_name[:-4]))
    sorted_video = (np.array(videos)[np.argsort(videos)])
    for img_s in sorted_video:
        sorted_imgs.append('{}.jpg'.format(img_s))
    return sorted_imgs

3. 绘制bbox可视化

def draw_yolox_predictions(img, bboxes, scores, bbclasses, confthre, classes_dict, gt=False, color=(0, 255, 0)):
    for i in range(len(bboxes)):
        box = bboxes[i]
        if not gt:
            score = scores[i]
            if score < confthre:
                continue
        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])

        cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
        if not gt:
            cls_id = int(bbclasses[i])
            cv2.putText(img, '{}:{:.1f}%'.format(0, score * 100), (x0, y0 - 3), cv2.FONT_HERSHEY_PLAIN, 1.3, color, thickness = 1)
    return img

4. 上下、左右、镜像翻转增强

#img size 720,1280,3
bboxes2, confs2  = predict(model2, img, size=IMG_SIZE, augment=AUGMENT)
bboxes3, confs3  = predict(model2, img[:,::-1,:], size=IMG_SIZE, augment=AUGMENT) #左右
bboxes4, confs4  = predict(model2, img[::-1,:,:], size=IMG_SIZE, augment=AUGMENT) #上下
bboxes5, confs5  = predict(model2, img[::-1,::-1,:], size=IMG_SIZE, augment=AUGMENT) #镜像
for bbox3, conf3 in zip(bboxes3, confs3):
     if conf3 >= 0.7:
         x3, y3, w3, h3 = bbox3
         bboxes_ensemble.append([1279-x3-w3, y3, 1279-x3, y3+h3])
         scores_ensemble.append(conf3)
    
for bbox4, conf4 in zip(bboxes4, confs4):
    if conf4 >= 0.7:
        x4, y4, w4, h4 = bbox4
        bboxes_ensemble.append([x4, 719-y4-h4, x4+w4, 719-y4])
        scores_ensemble.append(conf4)
        
for bbox5, conf5 in zip(bboxes5, confs5):
    if conf5 >= 0.75:
        x5, y5, w5, h5 = bbox5
        bboxes_ensemble.append([1279-x5-w5, 720-y5-h5, 1279-x5, 720-y5])
        scores_ensemble.append(conf5)

手动nms,输出最终结果

argsort = np.argsort(scores_ensemble)
argsort = argsort[::-1]
bboxes_sort = [bboxes_ensemble[p] for p in argsort]
scores_sort = [scores_ensemble[p] for p in argsort]

for a_i in range(len(bboxes_sort)):
    for b_j in range(a_i+1, len(bboxes_sort)):
        if scores_sort[b_j] > 0.15:
            bbox_A = bboxes_sort[a_i]
            bbox_B = bboxes_sort[b_j]
            x0_A, y0_A, x1_A,  y1_A = bbox_A
            x0_B, y0_B, x1_B,  y1_B = bbox_B
            iou = myiou(bbox_A, bbox_B)
            if iou > 0.3:
                scores_sort[b_j] = 0
                
take =  np.array(scores_sort) > 0.15
bboxes_nms = np.array(bboxes_sort)[take]
scores_nms = np.array(scores_sort)[take]
bboxes_nms = bboxes_nms.tolist()
scores_nms = scores_nms.tolist()

根据bbox裁剪正方形区域

#img size: 720,1280,3
def mycrop(box, enlarge):
    x0 = int(box[0])
    y0 = int(box[1])
    x1 = int(box[2])
    y1 = int(box[3])
    w = x1 - x0
    h = y1 - y0

    side = max(w, h)

    # enlarge = min(int(2*side), 719)
    enlarge_w = max(int((enlarge - w) / 2), 0)
    enlarge_h = max(int((enlarge - h) / 2), 0)

    xx0 = 0
    xx1 = 1279
    yy0 = 0
    yy1 = 719
    #超出边界
    if x0 < enlarge_w:
        xx0 = 0
        xx1 = enlarge
    if x1 > 1279 - enlarge_w:
        xx1 = 1279
        xx0 = 1279 - enlarge

    if y0 < enlarge_h:
        yy0 = 0
        yy1 = enlarge
    if y1 > 719 - enlarge_h:
        yy0 = 719 - enlarge
        yy1 = 719

    if x0 >= enlarge_w and x1 <= 1279 - enlarge_w:
        #larger
        xx0 = x0 - enlarge_w
        xx1 = x1 + enlarge_w

    if y0 >= enlarge_h and y1 <= 719 - enlarge_h:
        yy0 = y0 - enlarge_h
        yy1 = y1 + enlarge_h

    return xx0, xx1, yy0, yy1

计算相邻两帧相似度,判断是否属于同一个视频

def calculate(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0.0, 255.0])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0.0, 255.0])
    # plt.plot(hist1, color="r")
    # plt.plot(hist2, color="g")
    # plt.show()
    # 计算直方图的重合度
    degree = 0
    for i in range(len(hist1)):
        if hist1[i] != hist2[i]:
            degree = degree + (1 - abs(hist1[i] - hist2[i]) / max(hist1[i], hist2[i]))
        else:
            degree = degree + 1
    degree = degree / len(hist1)
    return degree

def similarity_img(img1, img2):
    sub_image1 = cv2.split(img1)
    sub_image2 = cv2.split(img2)
    sub_data = 0
    for im1, im2 in zip(sub_image1, sub_image2):
        sub_data += calculate(im1, im2)
    sub_data = sub_data / 3
    return sub_data