Spaces:
No application file
No application file
File size: 2,736 Bytes
e7c75bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import cv2
from utils import iou
from scipy import spatial
from darkflow.net.build import TFNet
import pandas as pd
def inference(file_name:str):
options = {'model': 'cfg/tiny-yolo-voc-3c.cfg',
'load': 3750,
'threshold': 0.1,
'gpu': 0.7}
tfnet = TFNet(options)
pred_bb = [] # predicted bounding box
pred_cls = [] # predicted class
pred_conf = [] # predicted class confidence
rbc = 0
wbc = 0
platelets = 0
cell = []
cls = []
conf = []
record = []
tl_ = []
br_ = []
iou_ = []
iou_value = 0
# tic = time.time()
image = cv2.imread(file_name)
output = tfnet.return_predict(image)
for prediction in output:
label = prediction['label']
confidence = prediction['confidence']
tl = (prediction['topleft']['x'], prediction['topleft']['y'])
br = (prediction['bottomright']['x'], prediction['bottomright']['y'])
if label == 'RBC' and confidence < .5:
continue
if label == 'WBC' and confidence < .25:
continue
if label == 'Platelets' and confidence < .25:
continue
# clearing up overlapped same platelets
if label == 'Platelets':
if record:
tree = spatial.cKDTree(record)
index = tree.query(tl)[1]
iou_value = iou(tl + br, tl_[index] + br_[index])
iou_.append(iou_value)
if iou_value > 0.1:
continue
record.append(tl)
tl_.append(tl)
br_.append(br)
center_x = int((tl[0] + br[0]) / 2)
center_y = int((tl[1] + br[1]) / 2)
center = (center_x, center_y)
if label == 'RBC':
color = (255, 0, 0)
rbc = rbc + 1
if label == 'WBC':
color = (0, 255, 0)
wbc = wbc + 1
if label == 'Platelets':
color = (0, 0, 255)
platelets = platelets + 1
radius = int((br[0] - tl[0]) / 2)
image = cv2.circle(image, center, radius, color, 2)
font = cv2.FONT_HERSHEY_COMPLEX
image = cv2.putText(image, label, (center_x - 15, center_y + 5), font, .5, color, 1)
cell.append([tl[0], tl[1], br[0], br[1]])
if label == 'RBC':
cls.append(0)
if label == 'WBC':
cls.append(1)
if label == 'Platelets':
cls.append(2)
conf.append(confidence)
# toc = time.time()
pred_bb.append(cell)
pred_cls.append(cls)
pred_conf.append(conf)
df = pd.DataFrame({"Cell type": ["rbc", "wbc", "platelets"], "Count": [rbc, wbc, platelets]})
return [image, df]
|