Spaces:
Runtime error
Runtime error
File size: 5,288 Bytes
fa14840 |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import sys
import cv2
import time
import numpy as np
import onnxruntime
print(onnxruntime.get_device())
lane_colors = [(249,65,68),(243,114,44),(248,150,30),(249,132,74),(249,199,79),(144,190,109),(77, 144, 142),(39, 125, 161)]
log_space = np.logspace(0,2, 50, base=1/10, endpoint=True)
class LSTR():
def __init__(self, model_path):
# Initialize model
self.model = self.initialize_model(model_path)
def __call__(self, image):
return self.detect_lanes(image)
def initialize_model(self, model_path):
opts = onnxruntime.SessionOptions()
opts.intra_op_num_threads = 16
self.session = onnxruntime.InferenceSession(model_path,sess_options=opts)
# Get model info
self.getModel_input_details()
self.getModel_output_details()
def detect_lanes(self, image):
input_tensor, mask_tensor = self.prepare_inputs(image)
outputs = self.inference(input_tensor, mask_tensor)
detected_lanes, good_lanes = self.process_output(outputs)
return detected_lanes, good_lanes
def prepare_inputs(self, img):
self.img_height, self.img_width, self.img_channels = img.shape
# Transform the image for inference
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img,(self.input_width, self.input_height))
# Scale input pixel values to -1 to 1
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
img = ((img/ 255.0 - mean) / std)
# img = img/ 255.0
img = img.transpose(2, 0, 1)
input_tensor = img[np.newaxis,:,:,:].astype(np.float32)
mask_tensor = np.zeros((1, 1, self.input_height, self.input_width), dtype=np.float32)
return input_tensor, mask_tensor
def inference(self, input_tensor, mask_tensor):
start = time.time()
outputs = self.session.run(self.output_names, {self.rgb_input_name: input_tensor,
self.mask_input_name: mask_tensor})
# print(time.time() - start)
return outputs
@staticmethod
def softmax(x):
"""Compute softmax values for each sets of scores in x."""
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum(axis=-1).T
def process_output(self, outputs):
pred_logits = outputs[0]
pred_curves = outputs[1]
# Filter good lanes based on the probability
prob = self.softmax(pred_logits)
good_detections = np.where(np.argmax(prob,axis=-1)==1)
pred_logits = pred_logits[good_detections]
pred_curves = pred_curves[good_detections]
lanes = []
for lane_data in pred_curves:
bounds = lane_data[:2]
k_2, f_2, m_2, n_1, b_2, b_3 = lane_data[2:]
# Calculate the points for the lane
y_norm = bounds[0]+log_space*(bounds[1]-bounds[0])
x_norm = (k_2 / (y_norm - f_2) ** 2 + m_2 / (y_norm - f_2) + n_1 + b_2 * y_norm - b_3)
lane_points = np.vstack((x_norm*self.img_width, y_norm*self.img_height)).astype(int)
lanes.append(lane_points)
self.lanes = lanes
self.good_lanes = good_detections[1]
return lanes, self.good_lanes
def getModel_input_details(self):
model_inputs = self.session.get_inputs()
self.rgb_input_name = self.session.get_inputs()[0].name
self.mask_input_name = self.session.get_inputs()[1].name
self.input_shape = self.session.get_inputs()[0].shape
self.input_height = self.input_shape[2]
self.input_width = self.input_shape[3]
def getModel_output_details(self):
model_outputs = self.session.get_outputs()
self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
# print(self.output_names)
def draw_lanes(self,input_img):
# Write the detected line points in the image
visualization_img = input_img.copy()
# Draw a mask for the current lane
right_lane = np.where(self.good_lanes==0)[0]
left_lane = np.where(self.good_lanes==5)[0]
if(len(left_lane) and len(right_lane)):
lane_segment_img = visualization_img.copy()
points = np.vstack((self.lanes[left_lane[0]].T,
np.flipud(self.lanes[right_lane[0]].T)))
cv2.fillConvexPoly(lane_segment_img, points, color =(0,191,255))
visualization_img = cv2.addWeighted(visualization_img, 0.7, lane_segment_img, 0.3, 0)
for lane_num,lane_points in zip(self.good_lanes, self.lanes):
for lane_point in lane_points.T:
cv2.circle(visualization_img, (lane_point[0],lane_point[1]), 3, lane_colors[lane_num], -1)
return visualization_img
if __name__ == '__main__':
model_path='../models/model_float32.onnx'
lane_detector = LSTR(model_path)
img = cv2.imread("../dog_road.jpg")
detected_lanes, lane_ids = lane_detector(img)
print(lane_ids)
lane_img = lane_detector.draw_lanes(img)
cv2.namedWindow("Detected lanes", cv2.WINDOW_NORMAL)
cv2.imshow("Detected lanes",lane_img)
cv2.imwrite("out.jpg", lane_img)
cv2.waitKey(0)
|