Take a look at my code, and see if there is something weird
def processa(self):
"""Create a TensorRT engine for ONNX-based YOLOv3-608 and run inference."""
# Try to load a previously generated YOLOv3-608 network graph in ONNX format:
onnx_file_path = "yolov3.onnx"
engine_file_path = "yolov3.trt"
# Download a dog image and save it to the following file path:
dev = cuda.Device(0) # 0 is your GPU number
ctx = dev.make_context()
with get_engine(
onnx_file_path, engine_file_path
) as engine, engine.create_execution_context() as context:
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
while True:
ativo, frame = self.imagem.read()
cv2.imwrite("frame.jpeg", frame)
input_image_path = "frame.jpeg"
# Two-dimensional tuple with the target network's (spatial) input resolution in HW ordered
input_resolution_yolov3_HW = (608, 608)
# Create a pre-processor object by specifying the required input resolution for YOLOv3
preprocessor = PreprocessYOLO(input_resolution_yolov3_HW)
# Load an image from the specified input path, and return it together with a pre-processed version
image_raw, image = preprocessor.process(input_image_path)
# Store the shape of the original input image in WH format, we will need it for later
shape_orig_WH = image_raw.size
# Output shapes expected by the post-processor
output_shapes = [(1, 255, 19, 19), (1, 255, 38, 38), (1, 255, 76, 76)]
# Do inference with TensorRT
trt_outputs = []
# Do inference
print("Running inference on image {}...".format(input_image_path))
# Set host input to the image. The common.do_inference function will copy the input to the GPU before executing.
inputs[0].host = image
trt_outputs = common.do_inference_v2(
context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream,
)
# Before doing post-processing, we need to reshape the outputs as the common.do_inference will give us flat arrays.
trt_outputs = [
output.reshape(shape)
for output, shape in zip(trt_outputs, output_shapes)
]
postprocessor_args = {
"yolo_masks": [
(6, 7, 8),
(3, 4, 5),
(0, 1, 2),
], # A list of 3 three-dimensional tuples for the YOLO masks
"yolo_anchors": [
(10, 13),
(16, 30),
(33, 23),
(30, 61),
(
62,
45,
), # A list of 9 two-dimensional tuples for the YOLO anchors
(59, 119),
(116, 90),
(156, 198),
(373, 326),
],
"obj_threshold": 0.6, # Threshold for object coverage, float value between 0 and 1
"nms_threshold": 0.5, # Threshold for non-max suppression algorithm, float value between 0 and 1
"yolo_input_resolution": input_resolution_yolov3_HW,
}
postprocessor = PostprocessYOLO(**postprocessor_args)
# Run the post-processing algorithms on the TensorRT outputs and get the bounding box details of detected objects
boxes, classes, scores = postprocessor.process(
trt_outputs, (shape_orig_WH)
)
objetos_captuados_frame = []
if boxes is not None:
for i in range(len(boxes)):
objeto_no_frame = {}
x = round(boxes[i][0])
y = round(boxes[i][1])
w = round(boxes[i][2])
h = round(boxes[i][3])
objeto_no_frame["object_id"] = classes[i]
objeto_no_frame["confianca"] = scores[i]
objeto_no_frame["topLeft"] = [x, y]
objeto_no_frame["bottomRight"] = [w, h]
objetos_captuados_frame.append(objeto_no_frame)
# Draw the bounding boxes onto the original input image and save it as a PNG file
image_raw = draw_bboxes(
image_raw, boxes, scores, classes, ALL_CATEGORIES
)
output_image_path = "dog_bboxes.png"
image_raw.save(output_image_path, "PNG")
# _, buffer = cv2.imencode(".png", obj_detected_img)
print(
"Saved image with bounding boxes of detected objects to {}.".format(
output_image_path
)
)
numpy_image = np.array(image_raw)
# convert to a openCV2 image, notice the COLOR_RGB2BGR which means that
# the color is converted from RGB to BGR format
image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
self.framecurrente = (
image # obj_detected_img = obj_detected_img.tobytes()
)
videredb.guardaFrame(
# obj_detected_img,
image_raw,
self.id_user,
time.time(),
objetos_captuados_frame,
)
ctx.pop()
del ctx