Spaces:
Sleeping
Sleeping
File size: 2,531 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 |
import cv2
import numpy as np
#import time
#video_path = 'D:/OfficeWork/VS_code_exp/exp/video_1.mp4'
#image_path = 'D:/OfficeWork/VS_code_exp/exp/test.jpg.jpg'
def load_model():
model= cv2.dnn.readNet(model='frozen_inference_graph.pb',
config='ssd_mobilenet_v2_coco_2018_03_29.pbtxt.txt',
framework='TensorFlow')
with open('object_detection_classes_coco.txt', 'r') as f:
class_names = f.read().split('\n')
COLORS = np.random.uniform(0, 255, size=(len(class_names), 3))
return model, class_names, COLORS
def load_img(img_path):
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 detect_objects(img, net):
blob = cv2.dnn.blobFromImage(img, size=(300, 300), mean=(104, 117, 123), swapRB=True)
net.setInput(blob)
outputs = net.forward()
#print (outputs)
return blob, outputs
def get_box_dimensions(outputs, height, width):
boxes = []
class_ids = []
for detect in outputs[0,0,:,:]:
scores = detect[2]
class_id = detect[1]
if scores > 0.3:
center_x = int(detect[0] * width)
center_y = int(detect[1] * height)
w = int(detect[5] * width)
h = int(detect[6] * height)
x = int((detect[3] * width))
y = int((detect[4] * height))
boxes.append([x, y, w, h])
class_ids.append(class_id)
return boxes, class_ids
def draw_labels(boxes, colors, class_ids, classes, img):
font = cv2.FONT_HERSHEY_PLAIN
model, classes, colors = load_model()
for i in range(len(boxes)):
x, y, w, h = boxes[i]
label = classes[int(class_ids[0])-1]
color = colors[i]
cv2.rectangle(img, (x,y), (w,h), color, 5)
cv2.putText(img, label, (x, y - 5), font, 5, color, 5)
return img
def image_detect(img_path):
model, classes, colors = load_model()
image, height, width, channels = load_img(img_path)
blob, outputs = detect_objects(image, model)
boxes, class_ids = get_box_dimensions(outputs, height, width)
image1 = draw_labels(boxes, colors, class_ids, classes, image)
return image1
#def start_video(video_path):
model, classes, colors = load_model()
cap = cv2.VideoCapture(video_path)
while True:
_, frame = cap.read()
height, width, channels = frame.shape
blob, outputs = detect_objects(frame, model)
boxes, class_ids = get_box_dimensions(outputs, height, width)
frame=draw_labels(boxes, colors, class_ids, classes, frame)
yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
cv2.destroyAllWindows() |