import cv2
import numpy as np
from rknnlite.api import RKNNLite
class yolop():
def __init__(self, confThreshold=0.25, nmsThreshold=0.5, objThreshold=0.45):
self.classes = ['car']
num_classes = len(self.classes)
anchors = [[3, 9, 5, 11, 4, 20], [
7, 18, 6, 39, 12, 31], [19, 50, 38, 81, 68, 157]]
self.nl = len(anchors)
self.na = len(anchors[0]) // 2
self.no = num_classes + 5
self.stride = np.array([8., 16., 32.])
self.anchor_grid = np.asarray(
anchors, dtype=np.float32).reshape(self.nl, -1, 2)
self.inpWidth = 640
self.inpHeight = 640
self.generate_grid()
self.net = RKNNLite()
self.net.load_rknn('./yolop.rknn')
self.net.init_runtime()
self.confThreshold = confThreshold
self.nmsThreshold = nmsThreshold
self.objThreshold = objThreshold
self.mean = np.array([0.485, 0.456, 0.406],
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.229, 0.224, 0.225],
dtype=np.float32).reshape(1, 1, 3)
self.keep_ratio = True
def generate_grid(self):
self.grid = [np.zeros(1)] * self.nl
self.length = []
self.areas = []
for i in range(self.nl):
h, w = int(self.inpHeight /
self.stride[i]), int(self.inpWidth/self.stride[i])
self.length.append(int(self.na * h * w))
self.areas.append(h*w)
if self.grid[i].shape[2:4] != (h, w):
self.grid[i] = self._make_grid(w, h)
def _make_grid(self, nx=20, ny=20):
xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)
def postprocess(self, frame, outs, newh, neww, padh, padw):
frameHeight = frame.shape[0]
frameWidth = frame.shape[1]
ratioh, ratiow = frameHeight / newh, frameWidth / neww
classIds = []
confidences = []
boxes = []
for detection in outs:
scores = detection[5:]
classId = np.argmax(scores)
confidence = scores[classId]
if confidence > self.confThreshold and detection[4] > self.objThreshold:
center_x = int((detection[0]-padw) * ratiow)
center_y = int((detection[1]-padh) * ratioh)
width = int(detection[2] * ratiow)
height = int(detection[3] * ratioh)
left = int(center_x - width / 2)
top = int(center_y - height / 2)
classIds.append(classId)
confidences.append(float(confidence) * detection[4])
boxes.append([left, top, width, height])
indices = cv2.dnn.NMSBoxes(
boxes, confidences, self.confThreshold, self.nmsThreshold)
for i in indices:
# i = i[0]
box = boxes[i]
left = box[0]
top = box[1]
width = box[2]
height = box[3]
frame = self.drawPred(
frame, classIds[i], confidences[i], left, top, left + width, top + height)
return frame
def drawPred(self, frame, classId, conf, left, top, right, bottom):
cv2.rectangle(frame, (left, top), (right, bottom),
(0, 0, 255), thickness=2)
label = '%.2f' % conf
label = '%s:%s' % (self.classes[classId], label)
labelSize, baseLine = cv2.getTextSize(
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
top = max(top, labelSize[1])
cv2.putText(frame, label, (left, top - 10),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2)
return frame
def resize_image(self, srcimg):
padh, padw, newh, neww = 0, 0, self.inpHeight, self.inpWidth
if self.keep_ratio and srcimg.shape[0] != srcimg.shape[1]:
hw_scale = srcimg.shape[0] / srcimg.shape[1]
if hw_scale > 1:
newh, neww = self.inpHeight, int(self.inpWidth / hw_scale)
img = cv2.resize(srcimg, (neww, newh),
interpolation=cv2.INTER_AREA)
padw = int((self.inpWidth - neww) * 0.5)
img = cv2.copyMakeBorder(img, 0, 0, padw, self.inpWidth - neww - padw, cv2.BORDER_CONSTANT,
value=0)
else:
newh, neww = int(self.inpHeight * hw_scale), self.inpWidth
img = cv2.resize(srcimg, (neww, newh),
interpolation=cv2.INTER_AREA)
padh = int((self.inpHeight - newh) * 0.5)
img = cv2.copyMakeBorder(
img, padh, self.inpHeight - newh - padh, 0, 0, cv2.BORDER_CONSTANT, value=0)
else:
img = cv2.resize(
srcimg, (self.inpWidth, self.inpHeight), interpolation=cv2.INTER_AREA)
return img, newh, neww, padh, padw
def _normalize(self, img):
img = img.astype(np.float32) / 255.0
img = (img - self.mean) / self.std
return img
def detect(self, srcimg):
img, newh, neww, padh, padw = self.resize_image(srcimg)
outs = self.net.inference(inputs=[img])
outs[1] = np.squeeze(outs[1], axis=-1)
outs[2] = np.squeeze(outs[2], axis=-1)
outimg = srcimg.copy()
drive_area_mask = outs[1][:, padh:(
self.inpHeight - padh), padw:(self.inpWidth - padw)]
seg_id = np.argmax(drive_area_mask, axis=0).astype(np.uint8)
seg_id = cv2.resize(
seg_id, (srcimg.shape[1], srcimg.shape[0]), interpolation=cv2.INTER_NEAREST)
outimg[seg_id == 1] = [0, 255, 0]
lane_line_mask = outs[2][:, padh:(
self.inpHeight - padh), padw:(self.inpWidth - padw)]
seg_id = np.argmax(lane_line_mask, axis=0).astype(np.uint8)
seg_id = cv2.resize(
seg_id, (srcimg.shape[1], srcimg.shape[0]), interpolation=cv2.INTER_NEAREST)
outimg[seg_id == 1] = [255, 0, 0]
det_out = outs[0]
row_ind = 0
for i in range(self.nl):
det_out[row_ind:row_ind+self.length[i], 0:2] = (
det_out[row_ind:row_ind+self.length[i], 0:2] * 2. - 0.5 + np.tile(self.grid[i], (self.na, 1))) * int(self.stride[i])
det_out[row_ind:row_ind+self.length[i], 2:4] = (
det_out[row_ind:row_ind+self.length[i], 2:4] * 2) ** 2 * np.repeat(self.anchor_grid[i], self.areas[i], axis=0)
row_ind += self.length[i]
outimg = self.postprocess(outimg, det_out, newh, neww, padh, padw)
return outimg
if __name__ == "__main__":
yolonet = yolop()
cap = cv2.VideoCapture('./test.mp4')
win = 'atk yolop det+seg demo'
cv2.namedWindow(win, 0)
cv2.setWindowProperty(win, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
while 1:
ret, srcimg = cap.read()
outimg = yolonet.detect(srcimg)
outimg = cv2.rotate(outimg, cv2.ROTATE_90_COUNTERCLOCKWISE)
cv2.imshow(win, outimg)
if cv2.waitKey(1) & 0xFF == ord('q'):
cap.release()
cv2.destroyAllWindows()
