python 图像iou_openCV——裁剪人的图片及iou自适应标签

import os

import cv2

import numpy as np

def IOU( box1, box2 ):

"""

: box1:[x1,y1,x2,y2]

: box2:[x1,y1,x2,y2]

# todo 1. 说明此函数的功能

2. 修改函数名称和返回值的名字，因为这不是真正的 IOU

"""

width1 = abs(box1[2] - box1[0])

height1 = abs(box1[1] - box1[3]) # 这里y1-y2是因为一般情况y1>y2，为了方便采用绝对值

width2 = abs(box2[2] - box2[0])

height2 = abs(box2[1] - box2[3])

x_max = max(box1[0],box1[2],box2[0],box2[2])

y_max = max(box1[1],box1[3],box2[1],box2[3])

x_min = min(box1[0],box1[2],box2[0],box2[2])

y_min = min(box1[1],box1[3],box2[1],box2[3])

iou_width = x_min + width1 + width2 - x_max

iou_height = y_min + height1 + height2 - y_max

if iou_width <= 0 or iou_height <= 0:

iou = 0

else:

iou_area = iou_width * iou_height # 交集的面积

# box1_area = width1 * height1

box2_area = width2 * height2

iou = iou_area / box2_area # 并集的面积

# if iou <= 0.6:

# iou = False

# else:

# iou = True

return iou

if __name__ == "__main__":

score_thr = 0.6

debug = False

remain_path = "../new/"

remain_txt_files = os.listdir(remain_path)

doc_dir = "E:/check/"

person_images = os.path.join(doc_dir, "images1/") # raw image

infer_person_txt = os.path.join(doc_dir, "/class.txt")

new_images = os.path.join(doc_dir, "new_image/")

new_labels = os.path.join(doc_dir, "new_label/")

res1 = dict() # key:str, image_id value:list of boxes, each box is a list of int of length 4

with open(infer_person_txt, 'r') as f:

for line in f:

split_line = line.split(" ")

if len(split_line) == 6:

img_id = split_line[0].strip()

conf = float(split_line[1])

box1 = [round(float(x)) for x in split_line[2:]]

if img_id not in res1:

res1[img_id] = list()

res1[img_id].append(box1)

res2 = dict() # key:str, image_id value:list of boxes, each box is a list of int of length 4

for file in remain_txt_files:

remain_txt = os.path.join(remain_path, file)

with open(remain_txt, 'r') as f:

for line in f:

split_line = line.split(" ")

if len(split_line) == 5:

img_id = split_line[0].strip()

box2 = [round(float(x)) for x in split_line[1:]]

if img_id not in res2:

res2[img_id] = list()

res2[img_id].append(box2) # 'conf': conf,

final_result = [] # each element: img_id, one_selected_worker_box

inter_set = set(res1.keys()) & set(res2.keys())

for k in inter_set:

res_person_infer_value = res1[k] # 从k中取出value,一串box,k图片对应的box

num_person = len(res_person_infer_value)

person_already_assigned = [0] * num_person # 0: not assigned; 1: assigned, correspond to res_main_infer_value

# print(main_already_assigned)

res2_remain_gt_value = res2[k]

# 对于 res2_remain_gt_value 中的每个 box(safebelt),

# 找到 res_main_infer_value 与其 IOU 值最大的 box(worker)

for sb_box in res2_remain_gt_value:

indexed_person_boxes = list(enumerate(res_person_infer_value)) # each element is like (idx, box)

filtered_indexed_person_boxes = \

[x for x in indexed_person_boxes if person_already_assigned[x[0]] == 0] # 筛选出未被 assign 的 workers

iou_score = np.array([IOU(x[1], sb_box) for x in filtered_indexed_person_boxes])

# print(iou_score)

max_index = np.argmax(iou_score) # [0]

max_score = np.max(iou_score)

if max_score > score_thr: # 大于设定值

max_score_person_index = filtered_indexed_person_boxes[max_index][0]

person_already_assigned[max_score_person_index] = 1

selected_person_box = filtered_indexed_person_boxes[max_index][1]

# ==== todo ===

# 求出 sb_box 相对 selected_worker_box 的 bounding box (绝对)坐标

box_relative = res2_remain_gt_value

# =============

tmp_triple = (k, selected_person_box, box_relative)

final_result.append(tmp_triple)

# 拿着 final_result 去抠图

# 1. 首先把 final_result

final_result_dict = dict()

for k, b1, b2 in final_result:

if k not in final_result_dict:

final_result_dict[k] = list()

final_result_dict[k].append((b1, b2))

# print(final_result_dict)

# 2. 拿着 final_result_dict 去 1.抠图 2. 写出相对person对属性的 bounding box 的绝对坐标

for k in final_result_dict.keys():

value_box = list(final_result_dict[k][0])

# print(value_box)

c, a, d, b = value_box[0]

old_loc = value_box[1][0]

# print(old_loc)

img_filename = os.path.join(person_images, k+'.jpg')

txt_name = k + ".txt"

txt_path = os.path.join(new_labels, txt_name)

img = cv2.imread(img_filename)

crop_img = img[a:b, c:d]

cv2.imwrite(new_images + k + '.jpg', crop_img)

img_w = d-c

img_h = b-a

if old_loc[0] < c:

old_loc[0] = c

if old_loc[1] < a:

old_loc[1] = a

if old_loc[2] > d:

old_loc[2] = d

if old_loc[3] > b:

old_loc[3] = b

x0 = old_loc[0] - c

y0 = old_loc[1] - a

x1 = old_loc[2] - c

y1 = old_loc[3] - a

# if x1 >= img_w: # ew_image的x轴

# x1 = img_w

# if y1 >= img_h: # new_image的y

# y1 = img_h

trans_x0 = round(float((x0+x1)/(2*img_w)), 4)

trans_y0 = round(float((y0+y1)/(2*img_h)), 4)

trans_x1 = round(float((x1-x0)/img_w), 4)

trans_y1 = round(float((y1-y0)/img_h), 4)

with open(txt_path, 'w') as new_file:

new_file.write('1' + ' ' + str(trans_x0) + ' ' + str(trans_y0) + ' ' + str(trans_x1)+ ' ' + str(trans_y1))

print('ok')

# == end your code here