Update main_tensorrt_gpu.py (#27)

This commit is contained in:
KendoClaw1 2022-08-24 19:41:21 +02:00 committed by GitHub
parent ad16aec34f
commit e1a9d45ff5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -97,6 +97,11 @@ def main():
im = cp.moveaxis(npImg, 3, 1) im = cp.moveaxis(npImg, 3, 1)
im = torch.from_numpy(cp.asnumpy(im)).to('cuda') im = torch.from_numpy(cp.asnumpy(im)).to('cuda')
#Converting to numpy for visuals
im0 = im[0].permute(1, 2, 0) * 255
im0 = im0.cpu().numpy().astype(np.uint8)
im0 = cv2.cvtColor(im0, cv2.COLOR_RGB2BGR) #Image has to be in BGR for visualization
# Detecting all the objects # Detecting all the objects
results = model(im) results = model(im)
@ -159,10 +164,10 @@ def main():
idx = 0 idx = 0
# draw the bounding box and label on the frame # draw the bounding box and label on the frame
label = "{}: {:.2f}%".format("Human", confidence * 100) label = "{}: {:.2f}%".format("Human", confidence * 100)
cv2.rectangle(npImg, (startX, startY), (endX, endY), cv2.rectangle(im0, (startX, startY), (endX, endY),
COLORS[idx], 2) COLORS[idx], 2)
y = startY - 15 if startY - 15 > 15 else startY + 15 y = startY - 15 if startY - 15 > 15 else startY + 15
cv2.putText(npImg, label, (startX, y), cv2.putText(im0, label, (startX, y),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLORS[idx], 2) cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLORS[idx], 2)
# Forced garbage cleanup every second # Forced garbage cleanup every second
@ -178,7 +183,7 @@ def main():
# See visually what the Aimbot sees # See visually what the Aimbot sees
if visuals: if visuals:
cv2.imshow('Live Feed', npImg) cv2.imshow('Live Feed', im0)
if (cv2.waitKey(1) & 0xFF) == ord('q'): if (cv2.waitKey(1) & 0xFF) == ord('q'):
exit() exit()
camera.stop() camera.stop()