--- a +++ b/predict_HRI.py @@ -0,0 +1,62 @@ +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)