predict_HRI.py

import cv2
from darkflow.net.build import TFNet

options = {'model': 'cfg/tiny-yolo-voc-3c.cfg',
           'load': 3750,
           'threshold': 0.15,
           'gpu': 0.7}

tfnet = TFNet(options)

C = []  # Center
R = []  # Radius
L = []  # Label

im_name = 'HRI001'
image = cv2.imread('data/' + im_name + '.jpg')

for h in range(0, 2592, 890):
    for w in range(0, 3872, 1290):
        im = image[h:h + 890, w:w + 1290]
        output = tfnet.return_predict(im)

        RBC = 0
        WBC = 0
        Platelets = 0

        for prediction in output:
            label = prediction['label']
            confidence = prediction['confidence']

            tl = (prediction['topleft']['x'], prediction['topleft']['y'])
            br = (prediction['bottomright']['x'], prediction['bottomright']['y'])

            height, width, _ = image.shape
            center_x = int((tl[0] + br[0]) / 2)
            center_y = int((tl[1] + br[1]) / 2)
            center = (center_x + w, center_y + h)
            radius = int((br[0] - tl[0]) / 2)

            C.append(center)
            R.append(radius)
            L.append(label)

record = []

for i in range(0, len(C)):
    center = C[i]
    radius = R[i]
    label = L[i]

    if label == 'RBC':
        color = (255, 0, 0)
    elif label == 'WBC':
        color = (0, 255, 0)
    elif label == 'Platelets':
        color = (0, 0, 255)

    image = cv2.circle(image, center, radius, color, 5)
    font = cv2.FONT_HERSHEY_COMPLEX
    image = cv2.putText(image, label, (center[0] - 30, center[1] + 10), font, 1, color, 2)

cv2.imwrite('output/' + im_name + 'out.jpg', image)