[Help]: 如何在自己的 Python 项目中使用 `tensorrt_yolo`

inm2002 commented 1 month ago

问题描述

你好，我希望将demo中的detect.py修改为输入opencv图像和输出opencv图像的形式，但是修改后无法正常进行推理。平均推理时间长达82ms，并且准确度下降，推理得到画出的检测框颜色不断发生变化。在此之前，我使用原版的detcet.py是可以正常推理的。我的tensorrt_yolo是自行编译的，Python版本为3.8，cuda版本为11.7 我有参考之前的关于输入为图像数组的issue，但我发现tensorrt_yolo包和他的好像不一样，所以就放弃了他的思路，自己写了。由于我无法正常稳定地上GitHub，所以暂时无法提供图像和视频的有关的不正常推理的资料，非常抱歉。

以下是我修改后的代码：

def detect_image(engine, mode, images):
    """
    YOLO Series Inference Script.
    """
    import cv2
    from rich.progress import track

    from tensorrt_yolo.infer import CpuTimer, DeployCGDet, DeployDet, GpuTimer, generate_labels_with_colors, visualize_detections

    labels = generate_labels_with_colors('labels_det.txt')
    is_obb = mode == 1

    model = DeployCGDet(engine, is_obb)
    results = model.predict(images)

    vis_image = visualize_detections(images, results, labels, is_obb)

    return cv2.cvtColor(vis_image, cv2.COLOR_RGB2BGR)

以下是我调用修改后的推理函数示例：

cv_img = cv2.cvtColor(np.array(img), cv2.COLOR_BGRA2RGB)

results = detect_image('best.engine', 0, cv_img)

希望能得到你的帮助，非常感谢

laugh12321 commented 1 month ago

@inm2002

首先，DeployCGDet.predict() 与 DeployDet.predict() 接收的参数就是 OpenCV 图像格式的数据（可以是一张图片也可以是多个图片组成的列表）但必须是RGB格式，不然会影响推理的精度。在demo.py 中已经给出了输入是列表的例子。 images = [cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB) for image_path in batch]
其次，generate_labels_with_colors 函数会生成以 (类别、颜色) 为元组的一组列表，其中的颜色是随机生成的。
最后，关于你说的每次推理颜色不断变化和推理很慢。是你调用方式的问题。你每次调用 detect_image 就会重新执行 generate_labels_with_colors 颜色当然不一样，而且每次调用你又重新去加载模型 model = DeployCGDet(engine, is_obb) 不慢才怪。每次调用只执行 model.predict 与 visualize_detections 就行，其它的只初始化一次，demo.py 写的很清楚。

以下是按照你的思路写的例子

from tensorrt_yolo.infer import DeployCGDet, generate_labels_with_colors, visualize_detections
import cv2
from pathlib import Path

def detect_image(model, labels, is_obb, image_file):
    """
    Detect objects in an image using a YOLO model and visualize the results.

    Parameters:
    - model: The YOLO model loaded with DeployCGDet.
    - labels: A list of labels for the detected objects.
    - is_obb: Boolean indicating whether to use Oriented Bounding Box (OBB) detection.
    - image_file: Path to the image file to be processed.

    Returns:
    - A BGR image with visualizations of detected objects.
    """
    try:
        # Read the image file and convert it to RGB color space
        image = cv2.cvtColor(cv2.imread(str(image_file)), cv2.COLOR_BGR2RGB)
        if image is None:
            raise ValueError(f"Failed to read image from {image_file}")

        # Perform object detection
        results = model.predict(image)

        # Visualize the detection results
        vis_image = visualize_detections(image, results, labels, is_obb)
        return cv2.cvtColor(vis_image[0], cv2.COLOR_RGB2BGR)
    except Exception as e:
        print(f"Error processing image {image_file}: {e}")
        return None

if __name__ == "__main__":
    # Set up the detection parameters
    is_obb = False
    labels = generate_labels_with_colors('labels_det.txt')
    model = DeployCGDet('model.engine', is_obb)

    # Get a list of image files to process
    image_files = list(Path("your/image/dir/").glob('*.png'))

    # Process each image file
    for image_file in image_files:
        vis_img = detect_image(model, labels, is_obb, image_file)
        if vis_img is not None:
            # Save or display the visualization
            cv2.imshow('Detection Result', vis_img)
            cv2.waitKey(0)  # Wait for a key press to continue
            cv2.destroyAllWindows()
            # Optionally save the image
            # cv2.imwrite(f"output_{image_file.stem}.png", vis_img)

inm2002 commented 1 month ago

非常感谢，学习到了，问题已经解决了。看了回答后才我发现我这个问题其实相当愚蠢，但大佬还是细心地给对我这种新手喂饭式的解答，真的非常感谢。

laugh12321 commented 1 month ago

@inm2002 没事，问题解决了就行。这个issue就关闭了。

laugh12321 / TensorRT-YOLO