|
| 1 | +import cv2 |
| 2 | +import matplotlib.pyplot as plt |
| 3 | + |
| 4 | +config_file = 'ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt' |
| 5 | +frozen_model = 'frozen_inference_graph.pb' |
| 6 | +model = cv2.dnn_DetectionModel(frozen_model, config_file) |
| 7 | + |
| 8 | +classLabels = [] |
| 9 | +file_name = 'labels.txt' |
| 10 | +with open(file_name, 'rt') as fpt: |
| 11 | + classLabels = fpt.read().rstrip('\n').split('\n') |
| 12 | + |
| 13 | +print(classLabels) |
| 14 | +print(len(classLabels)) |
| 15 | + |
| 16 | +model.setInputSize(320, 320) |
| 17 | +model.setInputScale(1.0/127.5) |
| 18 | +model.setInputMean((127.5, 127.5, 127.5)) |
| 19 | +model.setInputSwapRB(True) |
| 20 | + |
| 21 | +img = cv2.imread('boy.jpg') |
| 22 | +plt.imshow(img) |
| 23 | + |
| 24 | +ClassIndex, confidence, bbox = model.detect(img, confThreshold=0.5) |
| 25 | +print(ClassIndex) |
| 26 | + |
| 27 | +font_scale = 3 |
| 28 | +font = cv2.FONT_HERSHEY_PLAIN |
| 29 | +for ClassInd, conf, boxes in zip(ClassIndex.flatten(), confidence.flatten(), bbox): |
| 30 | + cv2.rectangle(img, boxes, (255, 0, 0), 2) |
| 31 | + cv2.putText(img, classLabels[ClassInd-1], (boxes[0]+10, boxes[1]+40), font, fontScale=font_scale, color=(0, 255, 0), thickness=3) |
| 32 | + |
| 33 | +plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) |
| 34 | +plt.show() |
| 35 | + |
| 36 | +# Webcam Object Detection |
| 37 | +cap = cv2.VideoCapture('1') |
| 38 | +if not cap.isOpened(): |
| 39 | + cap = cv2.VideoCapture(0) |
| 40 | +if not cap.isOpened(): |
| 41 | + raise IOError("Cannot open video") |
| 42 | + |
| 43 | +font_scale = 3 |
| 44 | +font = cv2.FONT_HERSHEY_PLAIN |
| 45 | + |
| 46 | +while True: |
| 47 | + ret, frame = cap.read() |
| 48 | + |
| 49 | + ClassIndex, confidence, bbox = model.detect(frame, confThreshold=0.55) |
| 50 | + print(ClassIndex) |
| 51 | + |
| 52 | + if (len(ClassIndex) != 0): |
| 53 | + for ClassInd, conf, boxes in zip(ClassIndex.flatten(), confidence.flatten(), bbox): |
| 54 | + if (ClassInd <= 80): |
| 55 | + cv2.rectangle(frame, boxes, (255, 0, 0), 2) |
| 56 | + cv2.putText(frame, classLabels[ClassInd - 1], (boxes[0] + 10, boxes[1] + 40), font, |
| 57 | + fontScale=font_scale, color=(0, 255, 0), thickness=3) |
| 58 | + |
| 59 | + cv2.imshow('Object Detection Tutorial', frame) |
| 60 | + |
| 61 | + if cv2.waitKey(2) & 0xFF == ord('q'): |
| 62 | + break |
| 63 | + |
| 64 | +cap.release() |
| 65 | +cv2.destroyAllWindows() |
0 commit comments