def test():
anchors = config.ANCHORS
transform = config.test_transforms
dataset = YOLODataset(
"COCO/train.csv",
"COCO/images/images/",
"COCO/labels/labels_new/",
S=[13, 26, 52],
anchors=anchors,
transform=transform,
)
S = [13, 26, 52]
scaled_anchors = torch.tensor(anchors) / (
1 / torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
loader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)
for x, y in loader:
boxes = []
for i in range(y[0].shape[1]):
anchor = scaled_anchors[i]
print(anchor.shape)
print(y[i].shape)
boxes += cells_to_bboxes(y[i],
is_preds=False,
S=y[i].shape[2],
anchors=anchor)[0]
boxes = nms(boxes,
iou_threshold=1,
threshold=0.7,
box_format="midpoint")
print(boxes)
plot_image(x[0].permute(1, 2, 0).to("cpu"), boxes)
def test(image, target):
gt_box = []
for y, anchor, recep_field in zip(target, anchors, config.GridSize):
# [batch_size, 3 * h * w, (score, x, y, w, h, num_classes)]
boxes = cells_to_bboxes(y, anchor, recep_field, is_preds=False)
index = np.where(boxes[..., 1:2] == 1)
boxes = boxes[index[:-1]]
gt_box += boxes[..., 2:]