Skip to content

Commit

Permalink
reformatted files to pep-8 standard
Browse files Browse the repository at this point in the history
  • Loading branch information
Vaibhav-Hariani committed May 21, 2024
1 parent 1e297fd commit 854eb21
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 109 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,21 @@
import sys
from cv_bridge import CvBridge
from yolo_interfaces.msg import Yolo



class Dist_To_Object(Node):
def __init__(self,model):
def __init__(self, model):
super().__init__('dist_to_obj_node')
self._bridge = CvBridge()
timer_period = 0.01 # seconds
self.timer = self.create_timer(timer_period, self.timer_callback)
self.img_publisher = self.create_publisher(Image, "/obj_detection_img", 10)
self.distToObj_publisher = self.create_publisher(Yolo, '/dist_to_obj',10)
self.img_publisher = self.create_publisher(
Image, "/obj_detection_img", 10)
self.distToObj_publisher = self.create_publisher(
Yolo, '/dist_to_obj', 10)
self.model = model
self.model.init_model()

def timer_callback(self):
msg = Yolo()

Expand All @@ -38,11 +41,10 @@ def timer_callback(self):

# Run object detection
self.model.objectDetection()
image = self.model.annotated_frame
image = self.model.annotated_frame
image = self._bridge.cv2_to_imgmsg(image, encoding="passthrough")
self.img_publisher.publish(image)


msg.dist = self.model.distance
msg.name = self.model.obj_name
msg.x = self.model.center_x
Expand All @@ -56,7 +58,6 @@ def timer_callback(self):
self.get_logger().info('Name to object: %s \n' % msg.name)
self.get_logger().info("\n ----------------------------------------------\n")



def main(args=None):
# rclpy.init(args=args) # Initialize ROS2 program
Expand All @@ -73,6 +74,7 @@ def main(args=None):
rclpy.spin(node)
node.destroy_node()
rclpy.shutdown()



if __name__ == '__main__':
main()
Original file line number Diff line number Diff line change
Expand Up @@ -4,141 +4,149 @@
import supervision as sv
import math
from ..ocr_submodule.ocr import OCR



class ObjectDetection():
def __init__(self,display=False):

def __init__(self, display=False):
self.display = display
init_params = sl.InitParameters()
self.cam = sl.Camera()

self.distance = 0.0
self.obj_x = 0.0
self.obj_y = 0.0
self.obj_name = ''

# We should tune these vals
init_params.coordinate_units = sl.UNIT.METER
init_params.coordinate_system = sl.COORDINATE_SYSTEM.RIGHT_HANDED_Z_UP_X_FWD
# For applications requiring long-range depth perception, they recommend setting depth_minimum_distance to 1m or more for improved performance.
# For applications requiring long-range depth perception, they
# recommend setting depth_minimum_distance to 1m or more for improved
# performance.
init_params.depth_minimum_distance = 1
init_params.depth_mode = sl.DEPTH_MODE.ULTRA
init_params.depth_maximum_distance = 2000
status = self.cam.open(init_params)
if status != sl.ERROR_CODE.SUCCESS:
print("Camera Open : "+repr(status)+". Exit program.")
print("Camera Open : " + repr(status) + ". Exit program.")
exit()
self.runtime = sl.RuntimeParameters()
# self.runtime.enable_fill_mode

def _display(self):
bounding_box_annotator = sv.BoundingBoxAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.TOP_CENTER)

bounding_box_img = bounding_box_annotator.annotate(
label_annotator = sv.LabelAnnotator(
text_position=sv.Position.TOP_CENTER)

bounding_box_img = bounding_box_annotator.annotate(
scene=self.img,
detections=self.detections
)

self.annotated_frame = label_annotator.annotate(
scene=bounding_box_img,
detections=self.detections
)
)

if not(len(self.detections) == 0):
if not (len(self.detections) == 0):
for detection in self.detections:
bbox = detection[0]
dist = self.distance_to_objects(bbox)
self.obj_name = self.model.names[detection[3]]


if "sign" in self.ocr.output:
self.obj_name = self.ocr.output
else:
self.obj_name = self.model.names[detection[3]]


distance_text = f"Distance: {dist:.2f} meters"
confidence_text = f"Confidence: {detection[2]:.2f}"
combined_text = f"{distance_text}\n{confidence_text}"

# Add the combined text to the annotated frame
cv2.putText(self.annotated_frame, combined_text, (int(bbox[0]), int(bbox[1] - 10)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
# this should be the default output when no detections occur
cv2.putText(self.annotated_frame, combined_text, (int(bbox[0]), int(
bbox[1] - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

# this should be the default output when no detections occur
else:
self.distance = math.nan
self.obj_name = "No Object"
self.obj_x = -1.00
self.obj_y = -1.00

cv2.imshow("IMAGE", self.annotated_frame)
cv2.imshow("IMAGE", self.annotated_frame)

def distance_to_objects(self,box):
try:
def distance_to_objects(self, box):
try:
# box = self.detections.xyxy[0]
center_x = int(box[0] + box[2])/2
center_y = int(box[1] + box[3])/2
err, pointCloudVal = self.point_cloud.get_value(round(center_x),round(center_y))
center_x = int(box[0] + box[2]) / 2
center_y = int(box[1] + box[3]) / 2
err, pointCloudVal = self.point_cloud.get_value(
round(center_x), round(center_y))
x_p = pointCloudVal[0]
y_p = pointCloudVal[1]
z_p = pointCloudVal[2]

#Find distance using Euclidean distance (in mm)
# Find distance using Euclidean distance (in mm)
distance = math.sqrt(x_p * x_p +
y_p * y_p +
z_p * z_p)
y_p * y_p +
z_p * z_p)

self.distance = distance
self.obj_x = x_p
self.obj_y = y_p

except IndexError:
except IndexError:
self.distance = math.nan
self.obj_name = "No Object"
self.obj_x = -1
self.obj_y = -1

return self.distance

def init_model(self):
image_size = self.cam.get_camera_information().camera_configuration.resolution

self.image_left_tmp = sl.Mat(self.cam.get_camera_information().camera_configuration.resolution.width, self.cam.get_camera_information().camera_configuration.resolution.height, sl.MAT_TYPE.U8_C4)
self.point_cloud = sl.Mat(self.cam.get_camera_information().camera_configuration.resolution.width, self.cam.get_camera_information().camera_configuration.resolution.height, sl.MAT_TYPE.U8_C4)

self.model=YOLO('best.pt')
self.image_left_tmp = sl.Mat(
self.cam.get_camera_information().camera_configuration.resolution.width,
self.cam.get_camera_information().camera_configuration.resolution.height,
sl.MAT_TYPE.U8_C4)
self.point_cloud = sl.Mat(
self.cam.get_camera_information().camera_configuration.resolution.width,
self.cam.get_camera_information().camera_configuration.resolution.height,
sl.MAT_TYPE.U8_C4)

self.model = YOLO('best.pt')
self.model.to('cpu')

self.ocr = OCR()

def close_cam(self):
cv2.destroyAllWindows()
self.cam.close()

def objectDetection(self):
#grab frames from ZED
# grab frames from ZED
err = self.cam.grab(self.runtime)
if err == sl.ERROR_CODE.SUCCESS: # Check that a new image is successfully acquired
if err == sl.ERROR_CODE.SUCCESS: # Check that a new image is successfully acquired
# cam.retrieve_image(imageZed, sl.VIEW.LEFT) # Retrieve left image
self.cam.retrieve_measure(self.point_cloud,sl.MEASURE.XYZRGBA)
self.cam.retrieve_measure(self.point_cloud, sl.MEASURE.XYZRGBA)
# cvImage = imageZed.get_data()
self.cam.retrieve_image(self.image_left_tmp, sl.VIEW.LEFT)
image_net = self.image_left_tmp.get_data()
self.img = cv2.cvtColor(image_net, cv2.COLOR_RGBA2RGB)
self.ocr(self.img)
cvPoint = self.point_cloud.get_data()


# # Crops the current image to view 1/3 of the frame
# # Crops the current image to view 1/3 of the frame
# center = self.img.shape
# x = center[1]/2
# x = center[1]/2
# y = center[0]/2
# t = 500
# self.img = self.img[int(y):int(y+t), int(x):int(x+t)]
results = self.model.predict(self.img)
self.detections = sv.Detections.from_ultralytics(results[0])

if self.display:
self._display()
self._display()
Loading

0 comments on commit 854eb21

Please sign in to comment.