def Supervision(img_data,out_13,out_26,out_52,turn):
out_13, out_26, out_52 =out_13.cpu().data, out_26.cpu().data, out_52.cpu().data
archor = tools.Archor()
box_13 = Select_Data(out_13, 32, archor[13])
box_26 = Select_Data(out_26, 16, archor[26])
box_52 = Select_Data(out_52, 8, archor[52])
boxes = torch.cat([box_13, box_26, box_52], dim=0)
boxes = boxes.numpy()
boxes = tools.NMS(boxes, 0.3)
#print(save_boxes)
dict_reverse = {0: '狗', 1: '人', 2: '羊驼', 3: '汽车', 4: '自行车', 5: '海豚', 6: '松鼠', 7: '马', 8: '马'}
img_data=np.array((img_data+0.5)*255,dtype=np.uint8)
img_data=np.transpose(img_data[0],(1,2,0))
img_back=Image.fromarray(img_data,"RGB")
img_draw=ImageDraw.Draw(img_back)
Font=ImageFont.truetype((r"F:\jkl\Yolo\simkai.ttf"),20)
#if turn>300:
plt.ion()
plt.clf()
for box in boxes:
img_draw.point((float(box[1]), float(box[2])), fill="red")
img_draw.rectangle((int(box[1] - box[3] / 2), int(box[2] - box[4] / 2), int(box[1] + box[3] / 2),
int(box[2] + box[4] / 2)), outline="red", width=2)
img_draw.text((int(box[1] - box[3] / 2), int(box[2] - box[4] / 2)),
text=dict_reverse[int(box[5])] + "{:.3f}".format(box[0]),
fill=(255, 0,0), font=Font)
plt.imshow(img_back)
img_back.save(r"F:\Reply\YoloV3\Vedio_Ready\{}.jpg".format(turn))
plt.pause(0.1)
plt.ioff()
def __getitem__(self, index):
feature_data = {}
line = self.datasets[0].split()
img_name = line[0]
line_data = []
img_open = Image.open(os.path.join(self.img_path, img_name))
img_data = tools.Trans_img(img_open)
for i in line[1:]:
line_data.append(float(i))
line_data = np.array(line_data)
boxes = np.split(line_data, len(line_data) // 5)
iou_dic = {}
for feature_size, W_H in tools.Archor().items():
feature_data[feature_size] = np.zeros(shape=(feature_size,
feature_size, 3,
5 + self.all_num))
# iou_dic[feature_size]=np.zeros(shape=(9,4))
# print(W_H)
for box in boxes:
cx, cy = float(box[1]), float(box[2])
cx_off, cx_index = math.modf(cx * feature_size / 416)
cy_off, cy_index = math.modf(cy * feature_size / 416)
w, h = int(box[3]), int(box[4])
for i, archor_area in enumerate(tools.Archor_Area().items()):
iou = tools.IOU_forlabel(box, W_H[i])
# print(iou)
t_w = w / W_H[i][0]
t_h = h / W_H[i][1]
one_hot = tools.One_Hot(int(self.all_num), int(box[0]))
iou_dic[iou] = [
iou, feature_size,
int(cy_index),
int(cx_index), i, box[0]
]
#print(np.array([iou,cx_off, cy_off, np.log(t_w), np.log(t_h),*one_hot]))
feature_data[feature_size][int(cy_index),
int(cx_index), i] = np.array([
iou, cx_off, cy_off,
np.log(t_w),
np.log(t_h), *one_hot
])
# print(iou_dic)
feature_data = tools.IOU_Deal(iou_dic, boxes, feature_data)
return img_data, torch.Tensor(feature_data[13]), torch.Tensor(
feature_data[26]), torch.Tensor(feature_data[52])
net = torch.load(net_path)
turn = 0
while True:
turn += 1
for i, (img_data, feature_13, feature_26,
feature_52) in enumerate(train_data):
img_data_cuda = img_data.to(device)
feature_13, feature_26, feature_52 = \
feature_13.to(device), feature_26.to(device), feature_52.to(device)
out_13, out_26, out_52 = net(img_data_cuda)
Supervision(img_data, out_13, out_26, out_52)
# print(out_13.shape,out_13)
# out_13, out_26, out_52=out_13.cpu().data, out_26.cpu().data, out_52.cpu().data
archor = tools.Archor()
loss_13 = Loss_design(out_13, feature_13, 0.9)
loss_26 = Loss_design(out_26, feature_26, 0.9)
loss_52 = Loss_design(out_52, feature_52, 0.9)
Loss = loss_13 + loss_26 + loss_52
optimizer.zero_grad()
Loss.backward()
optimizer.step()
print("{0}轮 -- Loss:{1}".format(turn, Loss.cpu().float()))
if turn % 20 == 0:
torch.save(net, net_path)
print("Save Successfully!")