Search code examples
opencvyolo

Can I change the batch size in a yolo detection model in OpenCV?

I am loading a yolo model with opencv in python with cv2.dnn_DetectionModel(cfg,weights) and then calling net.detect(img). I think I can get a speed-up per image using batches, but I don't see any support for batch size other than one.

Is it possible to set the batch size?

Solution

  • net.detect does not support batch size > 1.

    However, it's possible to do inference with batch size > 1 on darknet models with some extra work. Here is some partial code:

    net = cv2.dnn.readNetFromDarknet(cfg,weights)
net.setInputNames(["input"])
net.setInputShape("input",(batch_size,3,h,w))
blob = cv2.dnn.blobFromImages(image_list)
net.setInput(blob)
results = net.forward(net.getUnconnectedOutLayersNames())

    Now loop over all the images, and for each layer output in results, extract the boxes and confidences for this image each class, and having collected this info for every layer, pass this through cv2.dnn.NMSBoxes. This part is non-trivial, but doable.