Exemplo n.º 1
0
def batch_detection(network,
                    images,
                    class_names,
                    class_colors,
                    thresh=0.25,
                    hier_thresh=.5,
                    nms=.45,
                    batch_size=4):
    image_height, image_width, _ = check_batch_shape(images, batch_size)
    darknet_images = prepare_batch(images, network)
    batch_detections = darknet.network_predict_batch(network, darknet_images,
                                                     batch_size, image_width,
                                                     image_height, thresh,
                                                     hier_thresh, None, 0, 0)
    batch_predictions = []
    for idx in range(batch_size):
        num = batch_detections[idx].num
        detections = batch_detections[idx].dets
        if nms:
            darknet.do_nms_obj(detections, num, len(class_names), nms)
        predictions = darknet.remove_negatives(detections, class_names, num)
        images[idx] = darknet.draw_boxes(predictions, images[idx],
                                         class_colors)
        batch_predictions.append(predictions)
    darknet.free_batch_detections(batch_detections, batch_size)
    return images, batch_predictions
Exemplo n.º 2
0
def run_detection(output_splits, args):
    # Determine an optimal batch size
    if args.batch_size == None:
        args.batch_size = 1
        n_images = len(output_splits)
        while (n_images % (2 * args.batch_size) == 0):
            args.batch_size *= 2
        n_batches = int(n_images / args.batch_size)

    print('Optimal batch size: {}'.format(args.batch_size))
    print('Loading network.')

    # Load network
    random.seed(3)  # Deterministic bbox colours
    network, class_names, class_colors = darknet.load_network(
        args.config_file,
        args.data_file,
        args.weights,
        batch_size=args.batch_size)

    # Get network size
    n_w = darknet.network_width(network)
    n_h = darknet.network_height(network)

    # Split in batches
    global_predictions = []
    drawn_images = []
    for batch in range(n_batches):

        prev_time = time.time()  # Get batch calculation start time

        # Get images in the batch
        names_batch = output_splits[batch * args.batch_size:((batch + 1) *
                                                             args.batch_size)]
        images_batch = [cv2.imread(n) for n in names_batch]
        images_batch = [
            cv2.resize(n, (n_w, n_h), interpolation=cv2.INTER_LINEAR)
            for n in images_batch
        ]  # Ensure image resolution is the same as network's
        h, w, _ = check_batch_shape(images_batch, args.batch_size)
        dnet_images = darknet_images.prepare_batch(images_batch, network)

        # Detect
        batch_detections = darknet.network_predict_batch(
            network, dnet_images, args.batch_size, w, h, args.thresh,
            args.hier_thresh, None, 0, 0)

        # Process detections
        batch_predictions = []
        for idx in range(args.batch_size):
            num = batch_detections[idx].num
            detections = batch_detections[idx].dets
            if args.nms:
                darknet.do_nms_obj(detections, num, len(class_names), args.nms)
            predictions = darknet.remove_negatives(detections, class_names,
                                                   num)
            images_batch[idx] = darknet.draw_boxes(predictions,
                                                   images_batch[idx],
                                                   class_colors)
            batch_predictions.append(predictions)
        darknet.free_batch_detections(batch_detections, args.batch_size)

        # Store predictions and drawn images
        global_predictions.extend(batch_predictions)
        images_batch = [
            cv2.cvtColor(n, cv2.COLOR_BGR2RGB) for n in images_batch
        ]  # Ensure output is in RGB
        drawn_images.extend(images_batch)

        # Calculate FPS
        fps = int(args.batch_size / (time.time() - prev_time))
        print("FPS: {}".format(fps))

    return drawn_images, global_predictions, class_names
Exemplo n.º 3
0
        break

    arr = numpy.frombuffer(buf, dtype='uint8').reshape(
        (batch_size, height, width, 3))
    arr = arr.transpose((0, 3, 1, 2))
    arr = numpy.ascontiguousarray(arr.flat, dtype='float32') / 255.0
    darknet_images = arr.ctypes.data_as(darknet.POINTER(darknet.c_float))
    darknet_images = darknet.IMAGE(width, height, 3, darknet_images)
    raw_detections = darknet.network_predict_batch(net, darknet_images,
                                                   batch_size, width, height,
                                                   threshold, 0.5, None, 0, 0)
    detections = []
    for idx in range(batch_size):
        num = raw_detections[idx].num
        raw_dlist = raw_detections[idx].dets
        darknet.do_nms_obj(raw_dlist, num, len(class_names), 0.45)
        raw_dlist = darknet.remove_negatives(raw_dlist, class_names, num)
        dlist = []
        for cls, score, (cx, cy, w, h) in raw_dlist:
            dlist.append({
                'Category': cls,
                'Score': float(score),
                'Left': int(cx - w / 2),
                'Right': int(cx + w / 2),
                'Top': int(cy - h / 2),
                'Bottom': int(cy + h / 2),
            })
        detections.append(dlist)
    darknet.free_batch_detections(raw_detections, batch_size)
    print('json' + json.dumps(detections), flush=True)
Exemplo n.º 4
0
def darknet_batch_detection_rs_images(network,
                                      image_path,
                                      save_dir,
                                      patch_groups,
                                      patch_count,
                                      class_names,
                                      batch_size,
                                      thresh=0.25,
                                      hier_thresh=.5,
                                      nms=.45):
    '''
    # run batch detection of YOLO on an remote sensing image.
    :param network: a darknet network (already load the weight)
    :param image_path: the image patch.
    :param patch_groups: the group of images patch, each group has the same width and height.
    :param class_names:
    :param batch_size:
    :param thresh:
    :param hier_thresh:
    :param nms:
    :return:
    '''
    # read the entire image
    entire_img_data, nodata = raster_io.read_raster_all_bands_np(image_path)
    entire_img_data = entire_img_data.transpose(1, 2, 0)  # to opencv format
    # # RGB to BGR: Matplotlib image to OpenCV https://www.scivision.dev/numpy-image-bgr-to-rgb/
    # entire_img_data = entire_img_data[..., ::-1].copy()       # cancel RGB to BGR, since it make results worse, (maybe darknet already read images as RGB during training)
    entire_height, entire_width, band_num = entire_img_data.shape
    print("entire_height, entire_width, band_num", entire_height, entire_width,
          band_num)
    if band_num not in [1, 3]:
        raise ValueError('only accept one band or three band images')

    # class_colors={} #{'thaw_slump': (139, 140, 228)}  # class_colors should get when load network
    # for name in class_names:
    #     class_colors[name] = (0,0,0)

    patch_idx = 0
    for key in patch_groups.keys():
        patches_sameSize = patch_groups[key]

        batch_patches = [
            patches_sameSize[i * batch_size:(i + 1) * batch_size]
            for i in range((len(patches_sameSize) + batch_size - 1) //
                           batch_size)
        ]

        for a_batch_patch in batch_patches:
            t0 = time.time()
            input_batch_size = len(a_batch_patch)
            if input_batch_size != batch_size:
                if input_batch_size > batch_size:
                    raise ValueError(
                        'input_batch_size (%d) is greater than batch_size (%d)'
                        % (input_batch_size, batch_size))
                while (len(a_batch_patch) < batch_size):
                    a_batch_patch.append(
                        a_batch_patch[0]
                    )  # copy the first one to fit the batch size
            images = [
                copy_one_patch_image_data(patch, entire_img_data)
                for patch in a_batch_patch
            ]

            # for the batch prediction, it seems have to resize to the network size
            img_height, img_width, band_num = images[0].shape
            width = darknet.network_width(network)
            height = darknet.network_height(network)

            # darknet_images = prepare_batch(images, network)
            # prepare_batch
            darknet_images = []
            for image in images:
                # print('image shape',image.shape)
                image_resized = cv2.resize(image, (width, height),
                                           interpolation=cv2.INTER_LINEAR)
                if band_num == 1:
                    image_resized = np.expand_dims(
                        image_resized, axis=2)  # add one diminsion at the last
                # print('image_resized shape', image_resized.shape)
                custom_image = image_resized.transpose(2, 0, 1)
                darknet_images.append(custom_image)
            batch_array = np.concatenate(darknet_images, axis=0)
            batch_array = np.ascontiguousarray(batch_array.flat,
                                               dtype=np.float32) / 255.0
            darknet_images = batch_array.ctypes.data_as(
                darknet.POINTER(darknet.c_float))
            darknet_images = darknet.IMAGE(width, height, band_num,
                                           darknet_images)

            # prediction
            batch_detections = darknet.network_predict_batch(
                network, darknet_images, batch_size, img_height, img_width,
                thresh, hier_thresh, None, 0, 0)
            batch_predictions = []
            for idx in range(input_batch_size):
                num = batch_detections[idx].num
                detections = batch_detections[idx].dets
                if nms:
                    darknet.do_nms_obj(detections, num, len(class_names), nms)
                predictions = darknet.remove_negatives(detections, class_names,
                                                       num)
                # images[idx] = np.ascontiguousarray(images[idx])
                # images[idx] = darknet.draw_boxes(predictions, images[idx], class_colors)
                batch_predictions.append(predictions)
            darknet.free_batch_detections(batch_detections, batch_size)

            for idx, (patch, image, predictions) in enumerate(
                    zip(a_batch_patch, images, batch_predictions)):
                # remove the duplicated ones
                if idx >= input_batch_size:
                    break
                # save results
                save_res_json = os.path.join(save_dir, '%d.json' % patch_idx)
                # the confidence output by network_predict_batch is 0-1, convert to percent
                save_one_patch_detection_json(patch,
                                              predictions,
                                              class_names,
                                              save_res_json,
                                              b_percent=True)

                # cv2.imwrite('%d.png'%patch_idx, image)    # save for testing
                # for label, confidence, bbox in predictions:
                #     bbox = darknet.bbox2points(bbox)  # to [xmin, ymin, xmax, ymax]
                #     print(label, class_names.index(label), bbox, confidence)

                if patch_idx % 100 == 0:
                    print('saving %d patch, total: %d, cost %f second' %
                          (patch_idx, patch_count,
                           (time.time() - t0) / batch_size))

                patch_idx += 1