Diff of /predict_HRI.py [000000] .. [d34869]

Switch to side-by-side view

--- 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)