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