Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
#%% Define function to extract object coordinates if successful in detection
def where_is_it(frame, outputs):
frame_h = frame.shape[0]
frame_w = frame.shape[1]
bboxes, probs, class_ids = [], [], []
for preds in outputs: # different detection scales
hits = np.any(preds[:, 5:] > P_THRESH, axis=1) & (preds[:, 4] > OBJ_THRESH)
# Save prob and bbox coordinates if both objectness and probability pass respective thresholds
for i in np.where(hits)[0]:
pred = preds[i, :]
center_x = int(pred[0] * frame_w)
center_y = int(pred[1] * frame_h)
width = int(pred[2] * frame_w)
height = int(pred[3] * frame_h)
left = int(center_x - width / 2)
top = int(center_y - height / 2)
# Append all info
bboxes.append([left, top, width, height])
probs.append(float(np.max(pred[5:])))
class_ids.append(np.argmax(pred[5:]))
return bboxes, probs, class_ids
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment