-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathvehicle_cam_img_test.py
57 lines (50 loc) · 1.84 KB
/
vehicle_cam_img_test.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
import rospy
from pacmod_msgs.msg import PacmodCmd
import cv2
import torch
from cv_bridge import CvBridge, CvBridgeError
##FOr Webcam
class image_converter:
def __init__(self):
# self.image_pub = rospy.Publisher("image_topic_2",Image)
self.bridge = CvBridge()
# self.image_sub = rospy.Subscriber("image_topic",Image,self.callback)
self.image_sub = rospy.Subscriber("/mako_1/mako_1/image_raw", PacmodCmd, self.callback)
def callback(self,data):
global cv_image
try:
cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
except CvBridgeError as e:
print(e)
(rows,cols,channels) = cv_image.shape
if cols > 60 and rows > 60 :
cv2.circle(cv_image, (50,50), 10, 255)
cv2.imshow("Image window", cv_image)
cv2.waitKey(3)
# try:
# self.image_pub.publish(self.bridge.cv2_to_imgmsg(cv_image, "bgr8"))
# except CvBridgeError as e:
# print(e)
def main():
global cv_image
ic = image_converter()
rospy.init_node('Polaris GEM 2 Ex')
pub = rospy.Publisher("/pacmod/as_rx/turn_cmd", PacmodCmd, queue_size = 10)
# sub = rospy.Subscriber("/mako_1/mako_1/image_raw", PacmodCmd, callback)
#YOLO Network Here
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
count = 0
while not rospy.is_shutdown():
while not pub.get_num_connections() == 1:
pass
##What we need for image detection
results = model(cv_image) ## frame == image
temp = results.xyxy[0][:,5].tolist()
#image from the camera
if 0.0 in temp:
print("PERSON DETECTED")
brake(pub)
count += 1
print(count)
if __name__ == "__main__":
main()