Spaces:
No application file
No application file
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] | |