txt和xml互相转换

# gt_parse_yolo.py
import cv2
from gt_prob_class import *
import os


class FileName(object):
    def __init__(self, file_path):
        self.file_path = file_path
        self.full_name = os.path.split(file_path)[1]
        self.file_name = self.full_name.split('.')[0]


class DetectResult(object):
    def __init__(self, txt_path, img_path):
        self.datas = np.loadtxt(txt_path)
        self.len_datas = len(self.datas.shape)

        self.label = self.datas[:, 0]
        self.yolo_box = self.datas[:, [1, 2, 3, 4]]
        self.score = self.datas[:, -1]

        self.image = cv2.imread(img_path)
        self.h = self.image.shape[0]
        self.w = self.image.shape[1]
        self.b_box = self.yolo2xml()
        self.info = self.dict4prob()

    def yolo2xml(self):
        # yolo标注内容分别指类别、归一化后的中心点x坐标，归一化后的中心点y坐标，归一化后的目标框宽度w，归一化后的目标况高度h
        x0 = self.datas[:, 1] * self.w
        y0 = self.datas[:, 2] * self.h
        w0 = self.datas[:, 3] * self.w
        h0 = self.datas[:, 4] * self.h

        x1 = x0 - w0 / 2.0
        y1 = y0 - h0 / 2.0
        x2 = x0 + w0 / 2.0
        y2 = y0 + h0 / 2.0

        b_box = np.array([x1, y1, x2, y2]).T

        return b_box

    def dict4prob(self):
        info_1 = dict()
        info_1['bbox'] = self.b_box
        info_1['score'] = self.score
        info_1['class'] = self.label
        info_1['prob'] = self.score
        return info_1


def xml2yolo(size, label, box, score):
    print(box)  # x1, y1, x2, y2
    dw = 1. / size[1]
    dh = 1. / size[0]
    x = (box[:, 0] + box[:, 2]) / 2.0
    y = (box[:, 1] + box[:, 3]) / 2.0
    w = box[:, 2] - box[:, 0]
    h = box[:, 3] - box[:, 1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    new_box = np.column_stack((label, x, y, w, h, score))
    return new_box


def test_main(txt_path0, txt_path1, img_path0, save_name):
    # 文件路径信息
    # txt_path0 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01225.txt'
    # txt_path1 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01225.txt'
    # img_path0 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01225.jpg'
    img = cv2.imread(img_path0)

    # 得到字典
    detect_result1 = DetectResult(txt_path0, img_path0)
    detect_result2 = DetectResult(txt_path1, img_path0)

    # probEn计算
    keep, out_scores, out_boxes, out_class = test_vit(detect_result1.info, detect_result2.info)
    out_scores, out_boxes, out_class = np.array(out_scores), np.array(out_boxes), np.array(out_class)

    # 结果写入txt中
    yolo_content = xml2yolo(img.shape, out_class, out_boxes, out_scores)
    print('结果：')
    print('keep: \n', keep)
    print('out_scores: \n', out_scores)
    print('out_boxes: \n', np.array(out_boxes))
    print('out_class: \n', out_class)
    print('yolo_content: \n', yolo_content)

    save_path = ''
    np.savetxt(save_path + save_name, yolo_content, fmt='%f', delimiter=' ')


def process_all():
    uv_txt_dir = ''
    rgb_txt_dir = ''
    img_dir = ''
    file_names = os.listdir(uv_txt_dir)
    for file_name in file_names:
        uv_path = uv_txt_dir + file_name
        fm = FileName(file_name)
        rgb_path = rgb_txt_dir + file_name
        img_path = img_dir + fm.file_name + '.jpg'
        test_main(uv_path, rgb_path, img_path, file_name)


# process_all()

# 文件路径信息
txt_path0 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01225.txt'
txt_path1 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01227.txt'
img_path0 = r'D:\Graduate\dissertation\data_Kaist\gt_kaist\set00_V000_I01225.jpg'
save_name = '1.txt'
test_main(txt_path0, txt_path1, img_path0, save_name)

# gt_prob_class.py
import numpy as np
import torch


def weighted_box_fusion(bbox, score):
    weight = score / np.sum(score)
    out_bbox = np.array(bbox) * weight[:, None]
    out_bbox = np.sum(out_bbox, axis=0)
    return out_bbox


def prepare_data(info1, info2, info3='', method=None):
    out_dict = {}
    for key in info1.keys():
        if key != 'img_name':
            data1 = np.array(info1[key])
            data2 = np.array(info2[key])
            data_all = np.concatenate((data1, data2), axis=0)

            if info3:
                data3 = np.array(info3[key])
                data_all = np.concatenate((data_all, data3), axis=0)
            out_dict[key] = data_all
    return out_dict


def bayesian_fusion_single_class(match_score_vec, pred_class):
    match_score_vec.resize((1, 2))
    match_score_vec = match_score_vec.T
    scores = np.zeros((match_score_vec.shape[0], 2))
    scores[:, :1] = match_score_vec
    scores[:, -1] = 1 - np.sum(match_score_vec, axis=1)
    log_scores = np.log(scores)
    sum_logits = np.sum(log_scores, axis=0)
    exp_logits = np.exp(sum_logits)
    score_norm = exp_logits / np.sum(exp_logits)
    out_score = np.max(score_norm)
    out_class = np.argmax(score_norm)
    return out_score, out_class


def nms_bayesian(dict_collect, thresh, method):
    score_method, box_method = method
    classes = dict_collect['class']
    dets = dict_collect['bbox']
    scores = dict_collect['score']
    probs = dict_collect['prob']

    x1 = dets[:, 0] + classes * 640
    y1 = dets[:, 1] + classes * 512
    x2 = dets[:, 2] + classes * 640
    y2 = dets[:, 3] + classes * 512

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    out_classes = []
    match_scores = []
    match_bboxs = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        match = np.where(ovr > thresh)[0]
        match_ind = order[match + 1]

        match_prob = list(probs[match_ind])
        match_score = list(scores[match_ind])
        match_bbox = list(dets[match_ind][:, :4])
        original_prob = probs[i]
        original_score = scores[i].tolist()
        original_bbox = dets[i][:4]

        if len(match_score) > 0:
            match_score += [original_score]
            match_prob += [original_prob]
            match_bbox += [original_bbox]

            if score_method == "probEn":
                final_score, out_class = bayesian_fusion_single_class(np.asarray(match_prob), classes[i])
                out_classes.append(out_class)
            elif score_method == 'avg':
                final_score = np.mean(np.asarray(match_score))
                out_classes.append(classes[i])
            elif score_method == 'max':
                final_score = np.max(match_prob)
                out_classes.append(classes[i])

            if box_method == "v-avg":
                1
            elif box_method == 's-avg':
                final_bbox = weighted_box_fusion(match_bbox, match_score)
            elif box_method == 'avg':
                1
            elif box_method == 'argmax':
                max_score_id = np.argmax(match_score)
                final_bbox = match_bbox[max_score_id]

            match_scores.append(final_score)
            match_bboxs.append(final_bbox)
        else:
            match_scores.append(original_score)
            match_bboxs.append(original_bbox)
            out_classes.append(classes[i])

        order = order[inds + 1]

    assert len(keep) == len(match_scores)
    assert len(keep) == len(match_bboxs)
    assert len(keep) == len(out_classes)

    match_bboxs = match_bboxs
    match_scores = torch.Tensor(match_scores)
    match_classes = torch.Tensor(out_classes)

    return keep, match_scores, match_bboxs, match_classes


def test_vit(info_1, info_2):
    threshold = 0.5
    dict_collect = prepare_data(info_1, info_2)

    method = 'probEn', 's-avg'
    keep, out_scores, out_boxes, out_class = nms_bayesian(dict_collect, threshold, method=method)
    return keep, out_scores, out_boxes, out_class


def test_main():
    info_1 = {}
    # x1, y1, x2, y2(如果是yolo格式的话，需要乘以图像的宽高)
    info_1['bbox'] = [[295.4176330566406, 225.91998291015625, 332.5729675292969, 254.29959106445312],
                      [465.919921875, 226.74783325195312, 498.17022705078125, 251.9710235595703]]
    info_1['score'] = [0.9996856451034546, 0.9983585476875305]
    info_1['class'] = [0, 0]
    info_1['prob'] = info_1['score']

    info_2 = {}
    info_2['bbox'] = [[290.4176330566406, 220.91998291015625, 330.5729675292969, 250.29959106445312],
                      [463.919921875, 225.74783325195312, 494.17022705078125, 250.9710235595703]]
    info_2['score'] = [0.8996856451034546, 0.8883585476875305]
    info_2['class'] = [0, 0]
    info_2['prob'] = info_2['score']
    test_vit(info_1, info_2)


if __name__ == '__main__':
    test_main()

有许多其他自己的代码，做保存用。

互相转换体现在xml2yolo和yolo2xml中。