utils.py

import json

import torch
from torch import nn
from bs4 import BeautifulSoup
import os
from PIL import ImageDraw
import json
from matplotlib import pyplot as plt
class VOLoss(nn.Module):
    def __init__(self, l=(16, 16)):
        super(VOLoss, self).__init__()
        self.l = l
        self.siou = SIOULoss()
        self.aiou = AreaIOULoss()
        self.c = MLogClassifyLoss()

    def forward(self, input, target):
        """
        :param input:B,C(x,y,h,w,p,p_class1,p_class2...),H,W
        :param target:B,C(x,y,h,w,p,p_class1,p_class2...),H,W
        :return:
        """
        i_x_y_h_w = input[:, :4, :, :]
        i_xy = get_min_max(i_x_y_h_w, self.l)
        t_x_y_h_w = target[:, :4, :, :]
        t_xy = get_min_max(t_x_y_h_w, self.l)
        i_p = input[:, 4:5, :, :]
        t_p = target[:, 4:5, :, :]
        i_c = input[:, 5:, :, :]
        t_c = target[:, 5:, :, :]
        # iou=1.-self.iou(i_xy,t_xy)
        # loss_x_y_h_w=torch.mean(torch.sum(iou*t_p,dim=[-1,-2])/(torch.sum(t_p,dim=[-1,-2])+1e-5))
        loss_x_y_h_w = self.siou(i_xy, t_xy, mask=t_p)
        # loss_p=torch.mean(1.-torch.sum(t_p*i_p,dim=[-1,-2])/((torch.sum(i_p,dim=[-1,-2])+torch.sum(t_p,dim=[-1,-2]))/2+1e-5))
        loss_p = self.aiou(i_p, t_p)
        # loss_c=torch.mean(torch.mean(1.-torch.sum(i_c*t_c,dim=[-1,-2])/(torch.sum(t_c,dim=[-1,-2])+1e-5),dim=1))
        # loss_c=torch.mean(torch.mean(1.-torch.sum(t_c*i_c,dim=[-1,-2])/((torch.sum(i_c,dim=[-1,-2])+torch.sum(t_c,dim=[-1,-2]))/2+1e-5),dim=-1))
        # loss_c=torch.mean(torch.sum(-t_c*torch.log(i_c.clamp(0.001,0.999)),dim=[-1,-2])/(torch.sum(t_c,dim=[-1,-2])+1e-5))
        loss_c = self.c(i_c, t_c)
        return loss_x_y_h_w, loss_p, loss_c

    def iou(self, o: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        """B,C(x_min,x_max,y_min,y_max),H,W"""
        xl_min = label[:, 0:1]
        xl_max = label[:, 1:2]
        yl_min = label[:, 2:3]
        yl_max = label[:, 3:4]

        xo_min = o[:, 0:1]
        xo_max = o[:, 1:2]
        yo_min = o[:, 2:3]
        yo_max = o[:, 3:4]
        l = torch.min(torch.cat([xl_max, xo_max], dim=1), dim=1, keepdim=True).values - torch.max(
            torch.cat([xl_min, xo_min], dim=1), dim=1, keepdim=True).values
        h = torch.min(torch.cat([yl_max, yo_max], dim=1), dim=1, keepdim=True).values - torch.max(
            torch.cat([yl_min, yo_min], dim=1), dim=1, keepdim=True).values
        mj1 = (xl_max - xl_min) * (yl_max - yl_min)
        mj2 = (xo_max - xo_min) * (yo_max - yo_min)
        V = (-1 * torch.ones_like(l) + ((l.detach() > 0) * (h.detach() > 0)).float() * 2).detach()
        iou = V * torch.abs((l * h) / ((mj1 + mj2) / 2 + 1e-4))
        return iou


class AreaIOULoss(nn.Module):
    def forward(self, x, label):
        """
        :param x:shape(B,...)
        :param label:(B,...)
        :return:
        """
        shape = x.shape
        n_ = torch.sum(x * label, dim=[i for i in range(1, len(shape))]) if len(shape) > 1 else x * label
        d_1 = torch.sum(x, dim=[i for i in range(1, len(shape))]) if len(shape) > 1 else x
        d_2 = torch.sum(label, dim=[i for i in range(1, len(shape))]) if len(shape) > 1 else label
        loss = torch.mean(1. - n_ / ((d_1 + d_2) / 2 + 1e-5))
        return loss


class MLogClassifyLoss(nn.Module):
    def forward(self, x, label):
        """
        :param x:shape(B,C(class1,class2,...),...)
        :param label:shape(B,C(class1,class2,...),...)
        :return:
        """
        shape = x.shape
        n_ = torch.sum(-label * torch.log(x.clamp(0.001, 0.999)), dim=[i for i in range(2, len(shape))]) if len(
            shape) > 2 else -label * torch.log(x.clamp(0.001, 0.999))
        d_ = (torch.sum(label, dim=[i for i in range(2, len(shape))]) if len(shape) > 2 else label) + 1e-5
        loss = torch.mean(n_ / d_)
        return loss

class SIOULoss(nn.Module):
    def __init__(self):
        super(SIOULoss, self).__init__()

    def forward(self, x, label, mask=None):
        """
        :param x:shape(B,C[xmin,xmax,ymin,ymax],...)
        :param label:shape(B,C[xmin,xmax,ymin,ymax],...)
        :return:
        """
        shape = x.shape
        if mask == None:
            mask = torch.ones(shape[0], 1, shape[2:] if len(shape) > 2 else 1, dtype=torch.float, device=x.device)
        siou = self.get_siou(x, label)
        loss = torch.mean(torch.sum((1. - siou) * mask, dim=[i for i in range(1, len(shape))]) / (
                torch.sum(mask, dim=[i for i in range(1, len(shape))]) + 1e-5))
        return loss

    def get_siou(self, obj1: torch.Tensor, obj2: torch.Tensor) -> torch.Tensor:
        """
        :param obj1:shape(B,C[xmin,xmax,ymin,ymax],...)
        :param obj2:shape(B,C[xmin,xmax,ymin,ymax],...)
        :return:
        """
        x_obj1_min = obj1[:, 0:1]
        x_obj1_max = obj1[:, 1:2]
        y_obj1_min = obj1[:, 2:3]
        y_obj1_max = obj1[:, 3:4]

        x_obj2_min = obj2[:, 0:1]
        x_obj2_max = obj2[:, 1:2]
        y_obj2_min = obj2[:, 2:3]
        y_obj2_max = obj2[:, 3:4]
        L = torch.min(torch.cat([x_obj2_max, x_obj1_max], dim=1), dim=1, keepdim=True).values - torch.max(
            torch.cat([x_obj2_min, x_obj1_min], dim=1), dim=1, keepdim=True).values
        H = torch.min(torch.cat([y_obj2_max, y_obj1_max], dim=1), dim=1, keepdim=True).values - torch.max(
            torch.cat([y_obj2_min, y_obj1_min], dim=1), dim=1, keepdim=True).values
        S1 = (x_obj1_max - x_obj1_min) * (y_obj1_max - y_obj1_min)
        S2 = (x_obj2_max - x_obj2_min) * (y_obj2_max - y_obj2_min)
        P_M = (-1 * torch.ones_like(L) + ((L.detach() > 0) * (H.detach() > 0)).float() * 2).detach()
        siou = P_M * torch.abs((L * H) / ((S1 + S2) / 2 + 1e-4))
        return siou


def get_pad_(l=(8, 8)):
    l1 = torch.linspace(0, l[0] - 1, l[0])
    l2 = torch.linspace(0, l[1] - 1, l[1])
    x_p = (l1.view(1, 1, 1, -1).repeat([1, 1, l[1], 1])) / l[0]
    y_p = (l2.view(1, 1, -1, 1).repeat([1, 1, 1, l[0]])) / l[1]
    return x_p, y_p


# 通过txt内信息获取label
def get_label(cxyhw: list, num_classes, l=(8, 8)):
    """
    :param cxyhw:[[c,x,y,h,w]]
    :param l:
    :return:
    """
    label = torch.zeros(5 + num_classes, *l, dtype=torch.float)
    for i in cxyhw:
        class_ = i[0]
        x = i[1]
        y = i[2]
        h = i[3]
        w = i[4]
        b_x_l = 1 / (l[0] + 1e-4)
        b_y_l = 1 / (l[1] + 1e-4)
        n_b_x = int(x // b_x_l)
        n_b_y = int(y // b_y_l)
        x = (x % b_x_l) / (1 / (l[0] + 1e-4))
        y = (y % b_y_l) / (1 / (l[1] + 1e-4))
        label[0][n_b_y][n_b_x] = x
        label[1][n_b_y][n_b_x] = y
        label[2][n_b_y][n_b_x] = h
        label[3][n_b_y][n_b_x] = w
        label[4][n_b_y][n_b_x] = 1.
        label[class_ + 5][n_b_y][n_b_x] = 1.
    return label


# 获取xmax,ymax,xmin,ymin
def get_min_max(t, l=(16, 16)):
    """
    t.shape(B,C(x,y,h,w),H,W)
    """
    p_x, p_y = get_pad_(l)
    p_x = p_x.to(t.device)
    p_y = p_y.to(t.device)
    x = t[:, 0:1]
    y = t[:, 1:2]
    h = t[:, 2:3]
    w = t[:, 3:4]
    x = x / (l[0]) + p_x * (x != 0).float()
    y = y / (l[1]) + p_y * (y != 0).float()
    x_min = x - w / 2
    x_max = x + w / 2
    y_min = y - h / 2
    y_max = y + h / 2
    return torch.cat([x_min, x_max, y_min, y_max], dim=1)


# 计算siou
def iou(o: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
    """B,C(x_min,x_max,y_min,y_max),H,W"""
    xl_min = label[:, 0:1]
    xl_max = label[:, 1:2]
    yl_min = label[:, 2:3]
    yl_max = label[:, 3:4]

    xo_min = o[:, 0:1]
    xo_max = o[:, 1:2]
    yo_min = o[:, 2:3]
    yo_max = o[:, 3:4]
    l = torch.min(torch.cat([xl_max, xo_max], dim=1), dim=1, keepdim=True).values - torch.max(
        torch.cat([xl_min, xo_min], dim=1), dim=1, keepdim=True).values
    h = torch.min(torch.cat([yl_max, yo_max], dim=1), dim=1, keepdim=True).values - torch.max(
        torch.cat([yl_min, yo_min], dim=1), dim=1, keepdim=True).values
    mj1 = (xl_max - xl_min) * (yl_max - yl_min)
    mj2 = (xo_max - xo_min) * (yo_max - yo_min)
    V = (-1 * torch.ones_like(l) + ((l.detach() > 0) * (h.detach() > 0)).float() * 2).detach()
    iou = V * torch.abs((l * h) / ((mj1 + mj2) / 2 + 1e-4))
    return iou


# 将xml转为txt的函数，请忽略
def xml2txt(path, to_path, classes=('holothurian', 'echinus', 'scallop', 'starfish')):
    with open(path, "r", encoding="utf-8") as F:
        xml_str = F.read()
    b_obj = BeautifulSoup(xml_str, "lxml")
    f = b_obj.find_all("object")
    size = (float(b_obj.find("size").find("height").text), float(b_obj.find("size").find("width").text))
    all_msg = []
    for i in f:
        class_ = i.find("name").text
        bbox = i.find("bndbox")
        xmin = float(bbox.find("xmin").text)
        xmax = float(bbox.find("xmax").text)
        ymin = float(bbox.find("ymin").text)
        ymax = float(bbox.find("ymax").text)
        if class_ in classes:
            c_i = classes.index(class_)
            x = ((xmin + xmax) / 2) / size[1]
            y = ((ymin + ymax) / 2) / size[0]
            h = (ymax - ymin) / size[0]
            w = (xmax - xmin) / size[1]
            all_msg.append(" ".join(list(map(str, [c_i, x, y, h, w]))) + "\r")
    with open(to_path, "w", encoding="utf-8") as F:
        F.writelines(all_msg)
    # print(f)


# 文件夹内批量将xml转为txt
def fold_all_xml_txt(path, to_path):
    all_xml_files = os.listdir(path)
    for file in all_xml_files:
        t_f = file.replace("xml", "txt")
        full_x = os.path.join(path, file)
        full_t = os.path.join(to_path, t_f)
        xml2txt(full_x, full_t)


# 通过预测结果获取box
def get_bbox(x, p_=0.35, image_size=(512, 512), l=(16, 16)):
    """x.shape(c(x,y,h,w,p,class...),x,y)"""
    boxes = []
    C, H, W = x.shape
    x = torch.cat([get_min_max(x[:4].view(1, -1, H, W), l=l).view(-1, H, W), x[4:]], dim=0)
    x = x.view(x.shape[0], -1)
    x = x.permute(1, 0)
    bbox = [i for i in x if i[4].item() > p_]
    for b in bbox:
        boxes.append([b[0].item() * image_size[0], b[1].item() * image_size[0], b[2].item() * image_size[1],
                      b[3].item() * image_size[1], b[4].item(), torch.argmax(b[5:], dim=0).item()])
    return boxes


# 画框的函数
def draw_rectangle(image, bbox):
    image = image.copy()
    draw = ImageDraw.Draw(image)
    # 不同颜色框，让每个框更容易在图片中分辨
    color = [(0, 0, 0), (255, 255, 0), (255, 0, 0), (0, 0, 255)]
    for box in bbox:
        xmin, xmax, ymin, ymax, p, c_ = box
        draw.rectangle([xmin, ymin, xmax, ymax], outline=color[c_])
    return image


# NMS阈值抑制，将IOU大于某个数的框减少
def NMS(boxes, scores=0.6):
    keep_ = [True for box in boxes]
    for index in range(len(boxes)):
        if keep_[index] == False:
            continue
        xmin, xmax, ymin, ymax, p, c_ = boxes[index]
        for _index in range(len(boxes)):
            if _index == index or keep_[_index] == False:
                continue
            _xmin, _xmax, _ymin, _ymax, _p, _c_ = boxes[_index]
            a1 = (xmax - xmin) * (ymax - ymin)
            a2 = (_xmax - _xmin) * (_ymax - _ymin)
            w = min(xmax, _xmax) - max(xmin, _xmin)
            h = min(ymax, _ymax) - max(ymin, _ymin)
            if w > 0 and h > 0 and ((w * h) / ((a1 + a2) / 2 + 1e-4)) > scores:
                if a1 > a2:
                    keep_[_index] = False
                else:
                    keep_[index] = False
    r_boxes = []
    for index in range(len(boxes)):
        if keep_[index]:
            r_boxes.append(boxes[index])
    return r_boxes

#绘制loss曲线
def draw_loss(path):
    with open(path,"r",encoding="utf-8") as fp:
        all_data=json.load(fp)
    epchos=[i for i in range(len(all_data))]
    avg_loss=[i["avg_loss"] for i in all_data]
    bbox_loss=[i["bbox_loss"] for i in all_data]
    c_loss=[i["c_loss"] for i in all_data]
    p_loss=[i["p_loss"] for i in all_data]
    plt.plot(epchos,avg_loss,label="avg_loss",color="r")
    plt.plot(epchos,bbox_loss,label="bbox_loss",color="b")
    plt.plot(epchos,p_loss,label="p_loss",color="y")
    plt.plot(epchos,c_loss,label="c_loss",color="g")
    plt.legend(loc="upper right")  # 设置图例及图中文本显示
    plt.show()
if __name__ == '__main__':
    draw_loss(r"E:\pythonPro\deep_learning\CV_objDetection_0\log.json")