Ejemplo n.º 1
0
                                      num_workers=min(num_gpu * 8, 16))
    val_generator = data.DataLoader(val_dataset,
                                    batch_size=batch_size,
                                    shuffle=False,
                                    num_workers=min(num_gpu * 8, 16))

    # This is just to check that labels are correctly encoded.
    for pos in range(20):
        img1, target1 = train_dataset.__getitem__(pos)
        gt_color = (0, 0, 255)  #red

        # Get numpy image in opencv format
        np_img = preprocess.post_process_image(img1)
        # get bounding boxes: xyxy + class + confidence
        img_h, img_w, _ = np_img.shape
        target_bboxes = preprocess.decode_label(target1, img_h, img_w)

        draw_all_bboxes(np_img, target_bboxes, preprocess, gt_color,
                        "images_transformed/gt_{}.jpg".format(pos))

    import os
    os._exit(1)

    # Enable anomaly detection for debugging purpose
    torch.autograd.set_detect_anomaly(True)

    model = YOLO((channels, height, width), S, B, C)
    if num_gpu > 1:
        print("Let's use", num_gpu, "GPUs!")
        model = torch.nn.DataParallel(model)
    model = model.to(device)
Ejemplo n.º 2
0
        preds = model(batch_x)

        # compute loss
        test_loss = loss_func(preds, batch_y)
        accumulated_test_loss.append(test_loss)

        name = "test_results/images/test{}.jpg".format(it + 1)
        img = batch_x.clone().view((channels, height, width))
        pred = preds.clone().view((S, S, B * 5 + C))
        target = batch_y.clone().view((S, S, B * 5 + C))

        # Get numpy image in opencv format
        np_img = preprocess.post_process_image(img)
        # get bounding boxes: xyxy + class + confidence
        img_h, img_w, _ = np_img.shape
        target_bboxes = preprocess.decode_label(target, img_h, img_w)
        pred_bboxes = preprocess.decode_label(pred, img_h, img_w)

        draw_all_bboxes(np_img, target_bboxes, preprocess, (0, 0, 255), name)
        draw_all_bboxes(np_img, pred_bboxes, preprocess, (0, 255, 0), name)

        # write results to disk
        gt_file_name = "test_results/gt/test{}.txt".format(it + 1)
        write_result_to_disk(gt_file_name, target_bboxes, type="gt")
        pred_file_name = "test_results/pred/test{}.txt".format(it + 1)
        write_result_to_disk(pred_file_name, pred_bboxes, type="pred")

        it += 1

duration = time.time() - start_timestamp
# Epoch losses