def __init__(self,
model,
camera=None,
debug=False,
thresh=0.4,
last_network_callback_time=0.0):
self.debug = debug
if test_video_picture is not None:
self.camera = cv2.VideoCapture(test_video_picture)
elif self.debug:
self.camera = cv2.VideoCapture(camera)
self.camera.set(3, args.width)
self.camera.set(4, args.height)
self.model = model
self.detector = ObjectDetection(self.model, label_map=args.label)
self.thresh = float(thresh)
self.last_network_callback_time = last_network_callback_time
# Default to not running the network node on either camera, must push an enable/disable message from statemachine
# We do this so we can control what states the network is actually running in
self.run_network_front = False
self.run_network_bottom = False
def __init__(self, model, debug=False, fps=10.):
self.debug = debug
self.camera = MipiCamera(300, 300)
self.model = model
self.rate = float(1. / fps)
self.detector = ObjectDetection('./data/' + self.model)
class JetsonLiveObjectDetection():
def __init__(self,
model,
camera=None,
debug=False,
thresh=0.4,
last_network_callback_time=0.0):
self.debug = debug
if test_video_picture is not None:
self.camera = cv2.VideoCapture(test_video_picture)
elif self.debug:
self.camera = cv2.VideoCapture(camera)
self.camera.set(3, args.width)
self.camera.set(4, args.height)
self.model = model
self.detector = ObjectDetection(self.model, label_map=args.label)
self.thresh = float(thresh)
self.last_network_callback_time = last_network_callback_time
# Default to not running the network node on either camera, must push an enable/disable message from statemachine
# We do this so we can control what states the network is actually running in
self.run_network_front = False
self.run_network_bottom = False
def signal_handler(self, sig, frame):
''' Handles interrupt signals from OS.
Args:
signal: signal received from system
frame: state capture for diagnostic porpises.
'''
cv2.destroyAllWindows()
if args.debug or test_video_picture is not None:
self.camera.release()
self.detector.__del__()
sys.exit(0)
def _visualizeDetections(self, img, scores, boxes, classes,
num_detections):
''' Draws detections on images and returns a list of the detections name in string form
Args:
img: OpenCV Mat, The current OpenCV Image
scores: Float List, List of confidence values
boxes: 2D Float List, Bounding boxes for the objects
classes: Float List, List of the IDs for each detected object
num_detections: Float/Int, Number of detections found
Returns:
img: OpenCV Mat, Newly drawn-on image
detections: String List, Actual names of object detected in a list
'''
cols = img.shape[1]
rows = img.shape[0]
detections = []
for i in range(num_detections):
bbox = [float(p) for p in boxes[i]]
score = float(scores[i])
classId = int(classes[i])
if score > self.thresh:
detections.append(self.detector.labels[str(classId)])
x = int(bbox[1] * cols)
y = int(bbox[0] * rows)
right = int(bbox[3] * cols)
bottom = int(bbox[2] * rows)
thickness = int(4 * score)
cv2.rectangle(img, (x, y), (right, bottom), (125, 255, 21),
thickness=thickness)
cv2.putText(
img, self.detector.labels[str(classId)] + ': ' +
str(round(score, 3)), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1,
(255, 255, 255), 2)
return img, detections
def static_video(self):
''' Run the object detection on recorded video or images
'''
if test_video_picture is not None:
fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
names = test_video_picture.split('.')
if args.test_video is not None and not args.no_save_images:
out = cv2.VideoWriter(names[0] + '_output.' + names[1], fourcc,
30.0, (640, 480))
while (self.camera.isOpened()):
ret, img = self.camera.read()
if ret:
scores, boxes, classes, num_detections = self.detector.detect(
img)
img, new_detections = self._visualizeDetections(
img, scores, boxes, classes, num_detections)
if (args.show_video):
cv2.imshow(names[0], img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if args.test_video is not None and not args.no_save_images:
out.write(img)
elif args.test_picture is not None:
cv2.imwrite(names[0] + '_output.' + names[1], img)
print("Found objects: " + str(' '.join(new_detections)) +
".")
else:
break
if args.test_video is not None and not args.no_save_images:
out.release()
self.camera.release()
self.detector.__del__()
if (not args.no_save_images):
print("Output File written to " + names[0] + "_output." +
names[1])
exit()
def start(self):
''' Starts up the detector for static or live video.
'''
img_counter = 0
frame_counter = 0
if (not args.no_save_images):
script_directory = os.path.dirname(
os.path.realpath(__file__)) + '/'
save_dir = script_directory + 'saved_video/{}_{}/'.format(
self.model, datetime.datetime.now())
os.mkdir(save_dir)
print(
"Starting Live object detection, may take a few minutes to initialize..."
)
self.detector.initializeSession()
# Run the code with static video/picture testing:
if test_video_picture is not None:
self.static_video()
# Run the code in DEBUG mode. That is, with a local camera.
elif self.debug:
# Health Checks:
if not self.camera.isOpened():
print("Camera has failed to open")
exit(-1)
fps = self.camera.get(cv2.CAP_PROP_FPS)
# Main Programming Loop
while True:
curr_time = time.time()
ret, img = self.camera.read()
if (frame_counter >= fps / args.rate):
frame_counter = 0
scores, boxes, classes, num_detections = self.detector.detect(
img)
new_detections = None
img, new_detections = self._visualizeDetections(
img, scores, boxes, classes, num_detections)
print("Found objects: " + str(' '.join(new_detections)) +
".")
if (args.show_video):
cv2.imshow('Jetson Live Detection', img)
if ((img_counter % 3) == 0 and not args.no_save_images):
img_name = "{}opencv_frame_{}.jpg".format(
save_dir, int(curr_time))
cv2.imwrite(img_name, img)
img_counter = 0
img_counter += 1
# Publish ros-bridged images
if not args.no_ros:
img_msg = bridge.cv2_to_imgmsg(img)
front_img_pub.publish(img_msg)
detections_msg = Detections()
detections_msg.scores = scores
detections_msg.boxes = boxes.flatten()
detections_msg.classes = classes
detections_msg.detected = [num_detections]
front_detections_pub.publish(detections_msg)
print("Network running at: " +
str(1.0 / (time.time() - curr_time)) + " Hz.")
frame_counter += 1
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()
self.camera.release()
self.detector.__del__()
print("Exiting...")
return
# Run the code as a ROS node, pulls images on a topic, published them out on antoher
else:
rospy.Subscriber('front_raw_imgs',
Image,
self.run_network_node_front,
queue_size=1)
#rospy.Subscriber('bottom_raw_imgs', Image, self.run_network_node_bottom, queue_size=1)
rospy.Subscriber('enable_front_network', Bool,
self.enable_front_callback)
rospy.Subscriber('enable_bottom_network', Bool,
self.enable_bottom_callback)
rospy.spin()
def run_network_node_front(self, msg):
''' Runs network node on the recevial of an image from ROS
Args:
msg: ROS OpenCV Brige message, image to process on
'''
# Disable the front camera check
if not self.run_network_front:
return
if (time.time() - self.last_network_callback_time) <= args.rate:
return
bridge = cv_bridge.CvBridge()
img = bridge.imgmsg_to_cv2(msg)
scores, boxes, classes, num_detections = self.detector.detect(img)
new_detections = None
img, new_detections = self._visualizeDetections(
img, scores, boxes, classes, num_detections)
print("Found Front objects: " + str(' '.join(new_detections)) + ".")
img_msg = bridge.cv2_to_imgmsg(img)
front_img_pub.publish(img_msg)
detections_msg = Detections()
detections_msg.scores = scores
detections_msg.boxes = boxes.flatten()
detections_msg.classes = classes
detections_msg.detected = [num_detections]
front_detections_pub.publish(detections_msg)
self.last_network_callback_time = time.time()
def run_network_node_bottom(self, msg):
''' Runs network node on the recevial of an image from ROS
Args:
msg: ROS OpenCV Brige message, image to process on
'''
# Disable the bottom camera check
if not self.run_network_bottom:
return
if (time.time() - self.last_network_callback_time) <= args.rate:
return
bridge = cv_bridge.CvBridge()
img = bridge.imgmsg_to_cv2(msg)
scores, boxes, classes, num_detections = self.detector.detect(img)
new_detections = None
img, new_detections = self._visualizeDetections(
img, scores, boxes, classes, num_detections)
print("Found Bottom objects: " + str(' '.join(new_detections)) + ".")
img_msg = bridge.cv2_to_imgmsg(img)
#bottom_img_pub.publish(img_msg)
detections_msg = Detections()
detections_msg.scores = scores
detections_msg.boxes = boxes.flatten()
detections_msg.classes = classes
detections_msg.detected = [num_detections]
bottom_detections_pub.publish(detections_msg)
self.last_network_callback_time = time.time()
def enable_front_callback(self, msg):
self.run_network_front = msg.data
def enable_bottom_callback(self, msg):
self.run_network_bottom = msg.data
class JetsonLiveObjectDetection():
def __init__(self, model, debug=False, fps=10.):
self.debug = debug
self.camera = MipiCamera(300, 300)
self.model = model
self.rate = float(1. / fps)
self.detector = ObjectDetection('./data/' + self.model)
def _visualizeDetections(self, img, scores, boxes, classes,
num_detections):
cols = img.shape[1]
rows = img.shape[0]
detections = []
for i in range(num_detections):
bbox = [float(p) for p in boxes[i]]
score = float(scores[i])
classId = int(classes[i])
if score > 0.5:
x = int(bbox[1] * cols)
y = int(bbox[0] * rows)
right = int(bbox[3] * cols)
bottom = int(bbox[2] * rows)
thickness = int(4 * score)
cv2.rectangle(img, (x, y), (right, bottom), (125, 255, 21),
thickness=thickness)
detections.append(self.detector.labels[str(classId)])
print("Debug: Found objects: " + str(' '.join(detections)) + ".")
cv2.imshow('Jetson Live Detection', img)
def start(self):
print(
"Starting Live object detection, may take a few minutes to initialize..."
)
self.camera.startStreaming()
self.detector.initializeSession()
if not self.camera.isOpened():
print("Camera has failed to open")
exit(-1)
elif self.debug:
cv2.namedWindow("Jetson Live Detection", cv2.WINDOW_AUTOSIZE)
while True:
curr_time = time.time()
img = self.camera.getFrame()
scores, boxes, classes, num_detections = self.detector.detect(img)
if self.debug:
self._visualizeDetections(img, scores, boxes, classes,
num_detections)
print("Debug: Running at: " +
str(1.0 / (time.time() - curr_time)) + " Hz.")
if cv2.waitKey(1) == ord('q'):
break
# throttle to rate
capture_duration = time.time() - curr_time
sleep_time = self.rate - capture_duration
if sleep_time > 0:
time.sleep(sleep_time)
cv2.destroyAllWindows()
self.camera.__del__()
self.detector.__del__()
print("Exiting...")
return