'''
@Author: feizzhang
Created on: 21.01.2021
'''
import matplotlib.pyplot as plt
import numpy as np
import math
import argparse
import copy
import torch
from torch.autograd import Variable
INF = 100000.
'''
# --------------------------------
# Loss Function
# --------------------------------
'''
def loss_c(var_list, tgt_list, loss_type):
'''Compute the loss of N boxes and one GT that N boxes is to optimize
Args:
var_list: The variable need to optimize
[X, Y, W, H] X: tensor, N
tgt_list: The target to optimize
[X, Y, W, H] X: tensor, 1
Returns:
The loss value: tensor, N
'''
# --------------------------------
# Prepare the computation
# --------------------------------
X, Y, W, H = var_list
X0, Y0, W0, H0 = tgt_list
Xl, Yl, Xr, Yr = X - W / 2., Y - H / 2., X + W / 2., Y + H / 2.
X0l, Y0l, X0r, Y0r = X0 - W0 / 2., Y0 - H0 / 2., X0 + W0 / 2., Y0 + H0 / 2.
Xlmax, Xlmin = torch.max(Xl, X0l), torch.min(Xl, X0l)
Ylmax, Ylmin = torch.max(Yl, Y0l), torch.min(Yl, Y0l)
Xrmax, Xrmin = torch.max(Xr, X0r), torch.min(Xr, X0r)
Yrmax, Yrmin = torch.max(Yr, Y0r), torch.min(Yr, Y0r)
intersect_w = Xrmin - Xlmax
intersect_h = Yrmin - Ylmax
intersect_area = intersect_w.clamp(min=0) * intersect_h.clamp(min=0)
union_area = W * H + W0 * H0 - intersect_area
external_w = Xrmax - Xlmin
external_h = Yrmax - Ylmin
external_area = external_w * external_h
iou = (intersect_area / union_area).clamp(min=0)
# --------------------------------
# Compute the loss
# --------------------------------
if loss_type == "iou":
return 1 - iou
elif loss_type == "giou":
gitem = (external_area - union_area) / external_area
return 1 - iou + gitem
elif loss_type[:4] == "diou" or loss_type[:4] == "ciou":
distance_center = (X - X0) ** 2 + (Y - Y0) ** 2
distance_diagonal = external_w ** 2 + external_h ** 2 + 1e-16
ditem = distance_center / distance_diagonal
if loss_type[:4] == "diou":
if loss_type[4:] == "plus":
pass
else:
return 1 - iou + ditem
elif loss_type[:4] == "ciou":
v = (4 / math.pi ** 2) * torch.pow(torch.atan(W / H) - torch.atan(W0 / H0), 2)
with torch.no_grad():
alpha = v / (1 - iou + v)
citem = alpha * v
if loss_type[4:] == "plus":
pass
else:
return 1 - iou + ditem + citem
'''
# --------------------------------
# Define the training
# --------------------------------
'''
def simulation_engine(args):
# --------------------------------
# Initialize all
# --------------------------------
X = Variable(torch.FloatTensor([args.X]).cuda(), requires_grad=True)
Y = Variable(torch.FloatTensor([args.Y]).cuda(), requires_grad=True)
W = Variable(torch.FloatTensor([args.wh_base]).cuda(), requires_grad=True)
H = Variable(torch.FloatTensor([args.wh_base * args.ar]).cuda(), requires_grad=True)
X0 = torch.tensor([args.X0], dtype=torch.float16).cuda()
Y0 = torch.tensor([args.Y0], dtype=torch.float16).cuda()
W0 = torch.tensor([args.w0h0_base], dtype=torch.float16).cuda()
H0 = torch.tensor([args.w0h0_base * args.ar0], dtype=torch.float16).cuda()
var_list = [X, Y, W, H]
gt_list = [X0, Y0, W0, H0]
var_data = []
loss_data = []
# --------------------------------
# Define optimizer scheduler
# --------------------------------
# optimizer = torch.optim.Adam(var_list, lr=args.lr, weight_decay=0, betas=(0.9, 0.99))
optimizer = torch.optim.SGD(var_list,
lr=args.lr,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.gamma)
# --------------------------------
# Train for x iter
# --------------------------------
for iter in range(args.iter):
optimizer.zero_grad()
# --------------------------------
# 1) Save one iter variable
# --------------------------------
if iter % args.save_iter == 0:
# var_data.update({iter:copy.deepcopy(var_dict)})
temp = copy.deepcopy(var_list)
var_data.append({'X':temp[0],'Y':temp[1],'W':temp[2],'H':temp[3]})
# --------------------------------
# 2) Set the loss
# --------------------------------
loss = loss_c(var_list, gt_list, args.loss_type)
loss_data.append(loss)
# --------------------------------
# 3) Update grad once
# --------------------------------
loss.backward()
# --------------------------------
# 4) Update pred once
# --------------------------------
optimizer.step()
# --------------------------------
# 5) Update lr once
# --------------------------------
scheduler.step()
return var_data, loss_data
'''
# --------------------------------
# Simulate the boxes regression
# --------------------------------
'''
def main(args):
gt_ar_list = args.gt_ar_list
gt_wh_base_list = args.gt_wh_base_list
pred_ar_list = args.pred_ar_list
pred_wh_base_list = args.pred_wh_base_list
r = args.raduis * np.sqrt(np.random.rand(args.num_boxes))
thelta = 2 * np.pi * np.random.rand(args.num_boxes)
X = args.X0 + r * np.cos(thelta)
Y = args.Y0 + r * np.sin(thelta)
for iou_type in args.loss_type_list:
args.loss_type = iou_type
print("[^] {} begins!".format(iou_type))
num = 0
x_error = np.zeros(args.iter)
y_error = np.zeros(args.iter)
w_error = np.zeros(args.iter)
h_error = np.zeros(args.iter)
for ar, base in zip(gt_ar_list, gt_wh_base_list):
args.ar0 = ar
args.w0h0_base = base
for i in range(X.shape[0]):
args.X, args.Y = X[i], Y[i]
for ar, base in zip(pred_ar_list, pred_wh_base_list):
args.ar = ar
args.wh_base = base
num += 1
data, _ = simulation_engine(args)
for i, box in enumerate(data):
Xi = box['X'].detach().cpu().numpy()
Yi = box['Y'].detach().cpu().numpy()
Wi = box['W'].detach().cpu().numpy()
Hi = box['H'].detach().cpu().numpy()
x_error[i] += np.abs(Xi - args.X0)
y_error[i] += np.abs(Yi - args.Y0)
w_error[i] += np.abs(Wi - args.w0h0_base)
h_error[i] += np.abs(Hi - args.w0h0_base * args.ar0)
print(num, (x_error + y_error + w_error + h_error)[390:])
np2str = ''
error = x_error + y_error + w_error + h_error
for i in range(error.shape[0]):
if i == 0:
np2str = str(error[i])
if i > 0:
np2str = np2str + '_' + str(error[i])
if i == error.shape[0] - 1:
np2str = np2str + '\n'
with open(args.error_save_path, mode="a+") as f:
f.write(args.loss_type + '\n')
f.write(np2str)
print("[@] OVER sayolala!")
'''
# --------------------------------
# Plot the error
# --------------------------------
'''
def plot_error(args):
class results_tf(object):
def __init__(self):
super(results_tf, self).__init__()
storage = results_tf()
iou_type_list = []
with open(args.error_save_path, "r") as f:
lines = f.readlines()
for i in range(len(lines)//2):
iou_type = lines[2*i].strip()
iou_type_list.append(iou_type)
error = np.array(lines[2*i+1].strip().split('_')).astype(np.float32)
setattr(storage, iou_type, error.astype(int))
fig = plt.figure()
error1 = getattr(storage, iou_type_list[0])
x = np.arange(error1.shape[0])
for iou_i in iou_type_list:
errori = getattr(storage, iou_i)
plt.plot(x, errori, label=iou_i)
plt.legend(loc='upper right')
plt.ylim(0, 50000)
plt.show()
'''
# --------------------------------
# Visualize single box regression
# --------------------------------
'''
def visualize_box(args):
data, loss_rec = simulation_engine(args)
fig = plt.figure()
ax1 = fig.add_subplot(211)
for i, box in enumerate(data):
X = box['X'].cpu().detach().numpy()
Y = box['Y'].cpu().detach().numpy()
W = box['W'].cpu().detach().numpy()
H = box['H'].cpu().detach().numpy()
xl, yl, w, h = X - W / 2., Y - H / 2., W, H
color = "red"
alpha = 0.7
if i == len(data) - 1:
color = "yellow"
alpha = 1
ax1.add_patch(plt.Rectangle(xy=(xl, yl),
width=w,
height=h,
fill=False,
color=color,
alpha=alpha,
linewidth=1))
ax1.scatter(X, Y, s=15, alpha=alpha, color=color)
ax1.add_patch(plt.Rectangle(xy=(args.X0 - args.w0h0_base / 2., args.Y0 - args.w0h0_base * args.ar0 / 2.),
width=args.w0h0_base,
height=args.w0h0_base * args.ar0,
fill=False,
color="black",
alpha=0.5,
linewidth=2))
x_major_locator = plt.MultipleLocator(10)
y_major_locator = plt.MultipleLocator(10)
ax = plt.gca()
ax1.xaxis.set_major_locator(x_major_locator)
ax1.yaxis.set_major_locator(y_major_locator)
ax1.set_xlim(0, 10)
ax1.set_ylim(0, 10)
ax1.set_title(args.loss_type + " *** " + "config")
ax2 = fig.add_subplot(212)
losses = [loss.squeeze().cpu().detach().numpy() for loss in loss_rec]
x_ax = np.arange(0, len(losses), 1)
ax2.set_xlim(0, 400)
ax2.set_ylim(0, 2)
ax2.set_xlabel("epoch")
ax2.set_ylabel("losses")
ax2.plot(x_ax, losses)
plt.show()
'''
# --------------------------------
# Config
# --------------------------------
'''
parser = argparse.ArgumentParser(description="Set your own simulation hyperparameter")
# --------------------------------
# Initialize the training
# --------------------------------
parser.add_argument('--iter', type=int, default=400)
parser.add_argument('--save_iter', type=int, default=1, help="how iteration to save the data")
# --------------------------------
# Initialize the single prediction
# --------------------------------
parser.add_argument('--X', type=float, default=32.)
parser.add_argument('--Y', type=float, default=4.)
parser.add_argument('--ar', type=float, default=1.)
parser.add_argument('--wh_base', type=float, default=8.)
# --------------------------------
# Set the single GT
# --------------------------------
parser.add_argument('--X0', type=float, default=8.)
parser.add_argument('--Y0', type=float, default=8.)
parser.add_argument('--ar0', type=float, default=1.)
parser.add_argument('--w0h0_base', type=float, default=4.)
# --------------------------------
# Set the all the prediction
# and GT
# --------------------------------
parser.add_argument('--raduis', type=int, default=6)
parser.add_argument('--num_boxes', type=int, default=5000)
parser.add_argument('--error_save_path', type=str, default="./output.txt")
parser.add_argument('--gt_ar_list', type=list, default=[1., 4., 1 / 4, 2, 1 / 2, 3, 1 / 3])
parser.add_argument('--gt_wh_base_list', type=list, default=[2., 1, 4, 1.4, 2.8, 1.15, 3.46])
parser.add_argument('--pred_ar_list', type=list, default=[1., 3, 1 / 3, 2, 1 / 2, 4, 1 / 4])
parser.add_argument('--pred_wh_base_list', type=list, default=[3., 1.7, 5.1, 2.1, 4.2, 1.5, 6])
parser.add_argument('--loss_type_list', type=list, default=["diou", "ciou", "diouplus", "ciouplus"])
# --------------------------------
# Set the loss
# --------------------------------
parser.add_argument('--loss_type', type=str, default="ciouplus")
# --------------------------------
# Set optimizer & scheduler
# --------------------------------
parser.add_argument('--lr', type=float, default=5)
parser.add_argument('--gamma', type=float, default=0.3)
parser.add_argument('--weight_decay', type=float, default=0)
parser.add_argument('--milestones', type=list, default=[5, 60, 120, 200, 300])
# [5, 60, 120, 200, 300]
# [100, 180, 260, 320, 350] giou
# [40, 80, 260, 320, 350] diou [20, 80, 120, 240]
# [40, 80, 120, 180, 240, 300] ciou [20, 80, 120, 200]
if __name__ == '__main__':
args = parser.parse_args()
main(args)
