add config path for object detection

This commit is contained in:
shalenikol 2023-12-11 15:33:26 +03:00
parent eac306c479
commit 2db4d1eb67

View file

@ -4,12 +4,11 @@
ROS 2 program for 6D Pose Estimation
@shalenikol release 0.3
!!! version for testing 17.11.2023
"""
from typing import Optional
import os
import subprocess
import shutil
import json
import tempfile
@ -24,12 +23,16 @@ from rclpy.lifecycle import Publisher
from rclpy.lifecycle import State
from rclpy.lifecycle import TransitionCallbackReturn
from rclpy.timer import Timer
from rclpy.impl.logging_severity import LoggingSeverity
#from tf.transformations import quaternion_from_euler
from ament_index_python.packages import get_package_share_directory
from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import Pose, TransformStamped
from tf2_ros import TransformBroadcaster
from tf2_ros import TransformBroadcaster, TransformException
#from tf2_ros.buffer import Buffer
#from tf2_ros.transform_listener import TransformListener
from cv_bridge import CvBridge # Package to convert between ROS and OpenCV Images
import cv2 # OpenCV library
@ -393,7 +396,7 @@ class PoseEstimator(Node):
def yolo2megapose(self, res_json: str, path_to: Path) -> bool:
str_param = Path(res_json).read_text()
y = json.loads(str_param)[0]
conf = 0.75 # threshold of detection
conf = 0.5 # threshold of detection
found_coord = None
for detections in y["objects"]:
if detections["name"] == self.objName:
@ -441,12 +444,16 @@ class PoseEstimator(Node):
# object detection (YoloV4 - darknet)
self.get_logger().info(f"darknet: begin {self._image_cnt}")
result_json = str(self.objPath / "res.json")
#with open("log_darknet.txt", "w") as logf:
#self.get_logger().is_enabled_for(LoggingSeverity.WARN)
subprocess.run([darknet_bin, "detector", "test",
os.path.join(self.darknet_path, "yolov4_objs2.data"),
os.path.join(self.darknet_path, "yolov4_objs2.cfg"),
os.path.join(self.darknet_path, "yolov4.weights"),
"-dont_show", "-ext_output",
"-out", result_json, frame_im]
"yolov4_objs2.data",
"yolov4_objs2.cfg",
"yolov4.weights",
"-dont_show", "-ext_output", # -thresh 0.5
"-out", result_json, frame_im],
cwd = self.darknet_path #, stdout = subprocess.PIPE
)
detected = self.yolo2megapose(result_json, self.objPath)
self.get_logger().info(f"darknet: end {self._image_cnt}")