-
Notifications
You must be signed in to change notification settings - Fork 0
/
Object_detection.py
70 lines (61 loc) · 2.07 KB
/
Object_detection.py
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
import cv2
import numpy as np
import matplotlib.pyplot as plt
import numpy as np
yolo=cv2.dnn.readNet("2yolov3.weights","2yolov3.cfg")
classes=[]
with open("coco.names","r") as f:
classes=f.read().splitlines()
print(classes)
import cv2
cap=cv2.VideoCapture(0)
while True:
ret, frame =cap.read()
img=frame
# img=cv2.imread("test4.png")
# img=cv2.resize(img,(224,224))
height,width, _ =img.shape
# print(height,width)
blob= cv2.dnn.blobFromImage(img,1/255, (320,320), (0,0,0),swapRB=True,crop=False)
print(blob.shape)
# plt.imshow(img)
yolo.setInput(blob)
output_layers_name=yolo.getUnconnectedOutLayersNames()
layeroutput =yolo.forward(output_layers_name)
boxs=[]
confidences=[]
class_ids=[]
for output in layeroutput:
for detection in output:
score=detection[5:]
class_id=np.argmax(score)
confidence=score[class_id]
if confidence>0.7:
center_X=int(detection[0]*width)
center_y=int(detection[0]*height)
w=int(detection[0]*width)
h=int(detection[0]*height)
x=int(center_X-w/2)
y=int(center_y-h/2)
boxs.append([x,y,w,h])
confidences.append(float(confidence))
class_ids.append(class_id)
indexs=cv2.dnn.NMSBoxes(boxs,confidences,.5,.4)
font=cv2.FONT_HERSHEY_PLAIN
colors=np.random.uniform(0,255,size=(len(boxs),3))
if len(indexs) > 0:
for i in indexs.flatten():
x,y,w,h =boxs[i]
label=str(classes[class_ids[i]])
confi=str(round(confidences[i],2))
color=colors[i]
print(boxs[i])
cv2.rectangle(img,(x,y),(x+w,y+h),color,3)
cv2.putText(img,label+" "+confi,(x,y+20),font,2,(255,0,0),2)
cv2.imshow("window", img)
if cv2.waitKey(1) & 0xFF == ord('x'):
break
cv2.destroyAllWindows()
# plt.imshow(img)
# plt.show()
# cv2.imwrite("./img.jpg",img)