Spaces:
Sleeping
Sleeping
File size: 3,821 Bytes
fb26b5c |
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 133 134 |
import cv2
import numpy as np
#import argparse
#import time
#ideo_path = 'D:/OfficeWork/VS_code_exp/exp/video_1.mp4'
#image_path = 'D:\OfficeWork/VS_code_exp/exp/test.jpg.jpg'
#Load yolo
def load_yolo():
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
classes = []
with open("coco.names", "r") as f:
classes = [line.strip() for line in f.readlines()]
output_layers = [layer_name for layer_name in net.getUnconnectedOutLayersNames()]
colors = np.random.uniform(0, 255, size=(len(classes), 3))
return net, classes, colors, output_layers
def load_image(img_path):
# image loading
img = cv2.imread(img_path)
img = cv2.resize(img, None, fx=0.4, fy=0.4)
height, width, channels = img.shape
return img, height, width, channels
def start_webcam():
cap = cv2.VideoCapture(0)
return cap
def display_blob(blob):
'''
Three images each for RED, GREEN, BLUE channel
'''
for b in blob:
for n, imgb in enumerate(b):
cv2.imshow(str(n), imgb)
def detect_objects_yolo(img, net, outputLayers):
blob = cv2.dnn.blobFromImage(img, scalefactor=0.00392, size=(320, 320), mean=(0, 0, 0), swapRB=True, crop=False)
net.setInput(blob)
outputs = net.forward(outputLayers)
#output=np.ascontiguousarray(list(outputs))
#print(outputs)
#for i, out in enumerate(outputs):
# print(i, np.array(out).shape)
return blob, outputs
def get_box_dimensions_yolo(outputs, height, width):
boxes = []
confs = []
class_ids = []
for output in outputs:
for detect in output:
scores = detect[5:]
#print('detect', scores)
class_id = np.argmax(scores)
conf = scores[class_id]
if conf > 0.3:
center_x = int(detect[0] * width)
center_y = int(detect[1] * height)
w = int(detect[2] * width)
h = int(detect[3] * height)
x = int(center_x - w/2)
y = int(center_y - h / 2)
boxes.append([x, y, w, h])
#print(boxes)
confs.append(float(conf))
class_ids.append(class_id)
return boxes, confs, class_ids
def draw_labels_yolo(boxes, confs, colors, class_ids, classes, img):
indexes = cv2.dnn.NMSBoxes(boxes, confs, 0.5, 0.4)
font = cv2.FONT_HERSHEY_PLAIN
for i in range(len(boxes)):
if i in indexes:
x, y, w, h = boxes[i]
label = str(classes[class_ids[i]])
color = colors[i]
cv2.rectangle(img, (x,y), (x+w, y+h), color, 5)
cv2.putText(img, label, (x, y - 5), font, 5, color, 5)
return img
def image_detect_yolo(img_path):
model, classes, colors, output_layers = load_yolo()
image, height, width, channels = load_image(img_path)
blob, outputs = detect_objects_yolo(image, model, output_layers)
#print(outputs)
boxes, confs, class_ids = get_box_dimensions_yolo(outputs, height, width)
image=draw_labels_yolo(boxes, confs, colors, class_ids, classes, image)
return image
'''while True:
key = cv2.waitKey(1)
if key == 27:
break'''
#def webcam_detect():
model, classes, colors, output_layers = load_yolo()
cap = start_webcam()
while True:
_, frame = cap.read()
height, width, channels = frame.shape
blob, outputs = detect_objects(frame, model, output_layers)
boxes, confs, class_ids = get_box_dimensions(outputs, height, width)
draw_labels(boxes, confs, colors, class_ids, classes, frame)
key = cv2.waitKey(1)
if key == 27:
break
cap.release()
#def start_video_yolo(video_path):
model, classes, colors, output_layers = load_yolo()
cap = cv2.VideoCapture(video_path)
while True:
_, frame = cap.read()
height, width, channels = frame.shape
blob, outputs = detect_objects_yolo(frame, model, output_layers)
boxes, confs, class_ids = get_box_dimensions_yolo(outputs, height, width)
frame=draw_labels_yolo(boxes, confs, colors, class_ids, classes, frame)
yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
'''key = cv2.waitKey(1)
if key == 27 :
break
cap.release()'''
cv2.destroyAllWindows() |