Spaces:
Sleeping
Sleeping
File size: 4,359 Bytes
74cd228 |
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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
"""
Description:
"""
import cv2 as cv
import argparse
import sys
import numpy as np
import os.path
import matplotlib.pyplot as plt
def get_outputs_names(net):
"""Get the names of the output layers.
Args:
net: convolution object detection NN
"""
# Get the names of all the layers in the network
layersNames = net.getLayerNames()
# Get the names of the output layers, i.e. the layers with unconnected outputs
return [layersNames[i - 1] for i in net.getUnconnectedOutLayers()]
def draw_pred_box(class_id, conf, left, top, right, bottom, frame, classes, count=None):
"""Draw the predicted bounding box. Used in postProcess script.
Args:
- class_id: Indices for each of the object classes (e.g. herring)
- conf: confidence score for box
- left, top, right, bottom: frame indexes for corners of box
- frame: frame object
- classes: different class names
- count: current count of objects detected
"""
# Draw a bounding box.
cv.rectangle(frame, (left, top), (right, bottom), (255, 178, 50), 3)
label = "%.2f" % conf
count_label = "{}".format(count)
# Get the label for the class name and its confidence
if classes:
assert class_id < len(classes)
label = "%s:%s" % (classes[class_id], label)
# Display the label at the top of the bounding box
label_size, base_line = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
top = max(top, label_size[1])
cv.rectangle(
frame,
(left, top - round(1.5 * label_size[1])),
(left + round(1.5 * label_size[0]), top + base_line),
(255, 255, 255),
cv.FILLED,
)
cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 0), 1)
# Draw count in box
cv.putText(
img=frame,
text=count_label,
org=(right + 20, top + 50),
fontFace=cv.FONT_HERSHEY_SIMPLEX,
fontScale=1.75,
color=(255, 255, 255),
thickness=1,
)
def postprocess(frame, outs, conf_threshold, tracker, nms_threshold, classes):
"""Remove the bounding boxes with low confidence using non-maxima suppression.
Args:
- frame (frame object from cv2.VideoCapture): Image from video
- outs (nn layers): Output of the output layers from forward pass
"""
frame_height = frame.shape[0]
frame_width = frame.shape[1]
# Scan through all the bounding boxes output from the network and keep only the
# ones with high confidence scores. Assign the box's class label as the class with the highest score.
class_ids = []
confidences = []
boxes = []
counts = 0
for out in outs: # Scan through bounding boxes
for detection in out: # Scan through
scores = detection[5:]
class_id = np.argmax(scores) # class with highest score
confidence = scores[class_id] # confidence score for class
if confidence > conf_threshold:
print(detection)
center_x = int(detection[0] * frame_width)
center_y = int(detection[1] * frame_height)
width = int(detection[2] * frame_width)
height = int(detection[3] * frame_height)
left = int(center_x - width / 2)
top = int(center_y - height / 2)
class_ids.append(class_id)
confidences.append(float(confidence))
boxes.append([left, top, width, height])
counts += 1
# Perform non maximum suppression to eliminate redundant overlapping boxes with
# lower confidences.
indices = cv.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
print("BOXES", boxes, indices, class_ids)
max_counts = 0
for i in indices:
left, top, width, height = boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]
# temporary logic for counting all fish from detected boxes
if counts > max_counts:
max_counts = counts
draw_pred_box(
class_ids[i],
confidences[i],
left,
top,
left + width,
top + height,
frame=frame,
classes=classes,
count=counts,
)
print("counts", max_counts)
return max_counts
|