def show_predict(img_file_path):
img = cv2.imread(img_file_path)
color_mean = (104, 117, 123)
input_size = 300
transform = DataTransform(input_size, color_mean)
phase = "val"
img_tranformed, boxes, labels = transform(img, phase, "", "")
img_tensor = torch.from_numpy(img_tranformed[:,:,(2,1,0)]).permute(2,0,1)
net.eval()
input = img_tensor.unsqueeze(0) #(1, 3, 300, 300)
output = net(input)
plt.figure(figsize=(10, 10))
colors = [(255,0,0), (0,255,0), (0,0,255)]
font = cv2.FONT_HERSHEY_SIMPLEX
detections = output.data #(1, 21, 200, 5) 5: score, cx, cy, w, h
scale = torch.Tensor(img.shape[1::-1]).repeat(2)
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] >= 0.6:
score = detections[0, i, j, 0]
pt = (detections[0, i, j, 1:]*scale).cpu().numpy()
cv2.rectangle(img,
(int(pt[0]), int(pt[1])),
(int(pt[2]), int(pt[3])),
colors[i%3], 2
)
display_text = "%s: %.2f"%(classes[i-1], score)
cv2.putText(img, display_text, (int(pt[0]), int(pt[1])),
font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
j += 1
cv2.imshow("Result", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
"pottedplant", "sheep", "sofa", "train", "tvmonitor"
]
# prepare train, valid, annotation list
root_path = "./data/VOCdevkit/VOC2012/"
train_img_list, train_annotation_list, val_img_list, val_annotation_list = make_datapath_list(
root_path)
color_mean = (104, 117, 123)
input_size = 300
# transform = DataTransform(input_size, color_mean)
train_dataset = MyDataset(train_img_list,
train_annotation_list,
phase="train",
transform=DataTransform(input_size, color_mean),
anno_xml=Anno_xml(classes))
val_dataset = MyDataset(val_img_list,
val_annotation_list,
phase="val",
transform=DataTransform(input_size, color_mean),
anno_xml=Anno_xml(classes))
# print(train_dataset.__getitem__(1))
batch_size = 4
train_dataloader = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=my_collate_fn)
val_dataloader = data.DataLoader(val_dataset,
batch_size=batch_size,
if __name__ == "__main__":
classes = ['Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting',
'Cyclist', 'Tram', 'Misc', 'DontCare']
# prepare train, valid, annotation list
root_path = "../stereo_datasets/training"
train_img_list, train_annotation_list, val_img_list, val_annotation_list = make_datapath_list(root_path)
# prepare data transform
color_mean = (104, 117, 123)
input_size = 300
train_dataset = MyDataset(train_img_list, train_annotation_list, phase="train",
transform=DataTransform(input_size, color_mean), anno_xml=Anno_xml(classes))
val_dataset = MyDataset(val_img_list, val_annotation_list, phase="val",
transform=DataTransform(input_size, color_mean), anno_xml=Anno_xml(classes))
# print(len(train_dataset))
# print(train_dataset.__getitem__(1))
batch_size = 4
train_dataloader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=my_collate_fn)
val_dataloader = data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=my_collate_fn)
dataloader_dict = {
"train": train_dataloader,
"val": val_dataloader
}