生成验证集的结果,并检查验证集上的预测效果
[ptrblck] Add translation:
Generate the results of the validation set and check the prediction effect on the validation set
def boxes_to_lines(preds, image_id):
r = []
for bbox, label, score in zip(
preds[0]["boxes"].cpu().detach().numpy(),
preds[0]["labels"].cpu().detach().numpy(),
preds[0]["scores"].cpu().detach().numpy(),
):
# torchvision生成的bounding box格式为xyxy,需要转成xywh
xyxy = list(bbox)
xywh = [xyxy[0], xyxy[1], xyxy[2] - xyxy[0], xyxy[3]- xyxy[1]]
r.append(
{
"image_id": image_id,
"bbox": xywh,
"category_id": label,
"score": score,
}
)
return r
def output_result(image_paths):
result = []
model = get_model(num_classes=NUM_CLASS)
if os.path.exists(MODEL_PATH):
model.load_state_dict(torch.load(MODEL_PATH))
model.to(device)
model.eval()
for image_path in tqdm(image_paths):
img = Image.open(image_path)
image_id = os.path.basename(image_path).split(".")[0]
transform = T.Compose([T.ToTensor()]) # Defing PyTorch Transform
img_tensor = transform(img) # Apply the transform to the image
preds = model([img_tensor.to(device)])
result += boxes_to_lines(preds, image_id)
bboxes, labels = list(zip(*[
(bbox, label)
for bbox, label in zip(
preds[0]["boxes"].cpu().detach().numpy(),
preds[0]["labels"].cpu().detach().numpy()) ]))
img = draw_boxes(np.array(img), bboxes, labels)
plt.imshow(img)
return result
result = output_result(val_image_paths)