ROS/image_warp_node.py at master · cosplex01/ROS · GitHub

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
#!/user/bin/env python

import rospy
import cv2
import numpy as np
import os, rospkg
import json

from sensor_msgs.msg import CompressedImage
from cv_bridge import CvBridgeError

from utils import BEVTransform, draw_lane_img
from lane_detection import curve_learner

class IMGParser:
    def __init__(self):
        self.img_sub = rospy.Subscriber("/image_jpeg/compressed", CompressedImage, self.callback)
        self.img_wlane = None

    def callback(self, msg):
        try:
            np_arr = np.fromstring(msg.data, np.uint8)
            imb_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        except CvBridgeError as e:
            print(e)

        img_hsv = cv2.cvtColor(imb_bgr, cv2.COLOR_BGR2HSV)

        lower_wlane = np.array([0,0,190])
        upper_wlane = np.array([30,30,220])
        self.img_wlane = cv2.inRange(img_hsv, lower_wlane, upper_wlane)

if __name__ == '__main__':
    rp = rospkg.RosPack()

    currentPath = rp.get_path("lane_detection_example")
    with open(os.path.join(currentPath, 'sensor/sensor_params.json'), 'r') as fp:
        sensor_params = json.load(fp)
    params_cam = sensor_params["params_cam"]

    rospy.init_node('image_parser', anonymous=True)

    image_parser = IMGParser()
    bev_op = BEVTransform(params_cam=params_cam)

    rate = rospy.Rate(30)

    while not rospy.is_shutdown():
        if image_parser.img_wlane is not None:
           img_warp = bev_op.warp_bev_img(image_parser.img_wlane)
           lane_pts = bev_op.recon_lane_pts(image_parser.img_wlane)

           x_pred, y_pred_l, y_pred_r = curve_learner.fit_curve(lane_pts)

           xyl, xyr = bev_op.project_lane2img(x_pred, y_pred_l, y_pred_r)

           img_warp1 = draw_lane_img(img_warp, xyl[:,0].astype(np.int32),
                                    xyl[:,1].astype(np.int32),
                                    xyr[:,0].astype(np.int32),
                                    xyr[:,1].astype(np.int32))

           cv2.imshow("image Viewer", img_warp1)
           cv2.waitKey(1)

           rate.sleep()