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])
boxes = box_center.numpy()
save_boxes = tools.NMS(boxes, 0.3)
return save_boxes
if __name__ == '__main__':
dict_reverse = {0: '狗', 1: '人', 2: '羊驼', 3: '汽车', 4: '自行车', 5: '海豚', 6: '松鼠', 7: '马', 8: '猫'}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = torch.load(r"F:\jkl\Yolo_V3\Net_Save\Yolo_net_wen.pth")
net = net.to(device)
font=ImageFont.truetype((r"F:\jkl\Yolo\simkai.ttf"),25)
img_path = r"F:\Datasets_Dispose\Yolo_Datasets\Wen\1.jpg"
# for img_name in os.listdir(img_path):
img_open = Image.open(img_path)
img_draw = ImageDraw.Draw(img_open)
img_data = tools.Trans_img(img_open)
img_data = torch.unsqueeze(img_data, dim=0)
img_data = img_data.to(device)
out_13, out_26, out_52 = net(img_data)
out_13, out_26, out_52 = out_13.cpu().data, out_26.cpu().data, out_52.cpu().data
# print(out_13.shape,out_13)
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()
save_boxes = tools.NMS(boxes, 0.3, True)
print(save_boxes)
plt.ion()