Created
November 15, 2019 15:49
-
-
Save vinooniv/8082bfe067b202b9162c419eb716582a to your computer and use it in GitHub Desktop.
YOLO with OpenCV DNN module
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import argparse | |
import cv2 | |
import numpy as np | |
parser = argparse.ArgumentParser(add_help=False) | |
parser.add_argument("--image", default='samples/image.jpg', help="image for prediction") | |
parser.add_argument("--config", default='cfg/yolov3.cfg', help="YOLO config path") | |
parser.add_argument("--weights", default='yolov3.weights', help="YOLO weights path") | |
parser.add_argument("--names", default='data/coco.names', help="class names path") | |
args = parser.parse_args() | |
CONF_THRESH, NMS_THRESH = 0.5, 0.5 | |
# Load the network | |
net = cv2.dnn.readNetFromDarknet(args.config, args.weights) | |
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV) | |
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) | |
# Get the output layer from YOLO | |
layers = net.getLayerNames() | |
output_layers = [layers[i[0] - 1] for i in net.getUnconnectedOutLayers()] | |
# Read and convert the image to blob and perform forward pass to get the bounding boxes with their confidence scores | |
img = cv2.imread(args.image) | |
height, width = img.shape[:2] | |
blob = cv2.dnn.blobFromImage(img, 0.00392, (416, 416), swapRB=True, crop=False) | |
net.setInput(blob) | |
layer_outputs = net.forward(output_layers) | |
class_ids, confidences, b_boxes = [], [], [] | |
for output in layer_outputs: | |
for detection in output: | |
scores = detection[5:] | |
class_id = np.argmax(scores) | |
confidence = scores[class_id] | |
if confidence > CONF_THRESH: | |
center_x, center_y, w, h = (detection[0:4] * np.array([width, height, width, height])).astype('int') | |
x = int(center_x - w / 2) | |
y = int(center_y - h / 2) | |
b_boxes.append([x, y, int(w), int(h)]) | |
confidences.append(float(confidence)) | |
class_ids.append(int(class_id)) | |
# Perform non maximum suppression for the bounding boxes to filter overlapping and low confident bounding boxes | |
indices = cv2.dnn.NMSBoxes(b_boxes, confidences, CONF_THRESH, NMS_THRESH).flatten().tolist() | |
# Draw the filtered bounding boxes with their class to the image | |
with open(args.names, "r") as f: | |
classes = [line.strip() for line in f.readlines()] | |
colors = np.random.uniform(0, 255, size=(len(classes), 3)) | |
for index in indices: | |
x, y, w, h = b_boxes[index] | |
cv2.rectangle(img, (x, y), (x + w, y + h), colors[index], 2) | |
cv2.putText(img, classes[class_ids[index]], (x + 5, y + 20), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, colors[index], 2) | |
cv2.imshow("image", img) | |
cv2.waitKey(0) | |
cv2.destroyAllWindows() |
I am also facing same issue
any solution?
Probably one of the three below.
- Make sure you have yolov3 file in your model path. (In my case, the yolov3 file was a file that had not been downloaded)
- Check your image file. (It may not be in that path or it may be the wrong image.)
- The list of b_boxes may be empty.
AttributeError: 'tuple' object has no attribute 'flatten'
Can I have some tips how to solve it?
Following up can you please help if anyone solved?
@darius-luca-tech did you find any solution to the same?
TIA.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
@darius-luca-tech I am getting the same error.
Please Help!!