YOLOv5怎么以txt或json格式輸出預測結果

YOLOv5怎么以txt或json格式輸出預測結果

YOLOv5 是一個流行的目標檢測模型，廣泛應用于計算機視覺任務中。在實際應用中，我們通常需要將模型的預測結果以特定的格式輸出，以便后續處理或分析。本文將詳細介紹如何使用 YOLOv5 以 txt 或 json 格式輸出預測結果。

1. YOLOv5 預測結果的基本格式

YOLOv5 的預測結果通常包含以下信息：

• 目標類別（class）
• 目標邊界框（bounding box）
• 目標置信度（confidence）

這些信息可以通過 YOLOv5 的推理腳本直接輸出到控制臺，或者保存到文件中。

2. 以 txt 格式輸出預測結果

2.1 修改 YOLOv5 的推理腳本

YOLOv5 的推理腳本 detect.py 默認會將預測結果保存為圖像文件，并在圖像上繪制檢測框。為了將預測結果保存為 txt 文件，我們需要對 detect.py 進行一些修改。

首先，找到 detect.py 中保存結果的部分代碼：

# Save results (image with detections)
if save_img:
    if dataset.mode == 'image':
        cv2.imwrite(save_path, im0)
    else:  # 'video' or 'stream'
        if vid_path != save_path:  # new video
            vid_path = save_path
            if isinstance(vid_writer, cv2.VideoWriter):
                vid_writer.release()  # release previous video writer
            fps = vid_cap.get(cv2.CAP_PROP_FPS)
            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        vid_writer.write(im0)

在這段代碼之后，我們可以添加保存 txt 文件的邏輯。假設我們希望每個圖像對應一個 txt 文件，文件中每一行表示一個檢測到的目標，格式為 class x_center y_center width height confidence。

# Save results to txt file
if save_txt:
    txt_path = os.path.splitext(save_path)[0] + '.txt'
    with open(txt_path, 'w') as f:
        for *xyxy, conf, cls in reversed(det):
            x_center = (xyxy[0] + xyxy[2]) / 2 / im0.shape[1]
            y_center = (xyxy[1] + xyxy[3]) / 2 / im0.shape[0]
            width = (xyxy[2] - xyxy[0]) / im0.shape[1]
            height = (xyxy[3] - xyxy[1]) / im0.shape[0]
            f.write(f'{int(cls)} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {conf:.6f}\n')

2.2 運行推理腳本

在修改完 detect.py 后，我們可以通過以下命令運行推理腳本，并指定保存 txt 文件的路徑：

python detect.py --source path/to/images --save-txt --output path/to/output

其中，--source 指定輸入圖像或視頻的路徑，--save-txt 表示保存 txt 文件，--output 指定輸出路徑。

3. 以 json 格式輸出預測結果

3.1 修改 YOLOv5 的推理腳本

與 txt 格式類似，我們可以通過修改 detect.py 來保存 json 格式的預測結果。json 格式通常更適合存儲結構化數據，并且可以包含更多的信息。

在 detect.py 中，找到保存結果的部分代碼，并添加保存 json 文件的邏輯：

import json

# Save results to json file
if save_json:
    json_path = os.path.splitext(save_path)[0] + '.json'
    results = []
    for *xyxy, conf, cls in reversed(det):
        x_center = (xyxy[0] + xyxy[2]) / 2 / im0.shape[1]
        y_center = (xyxy[1] + xyxy[3]) / 2 / im0.shape[0]
        width = (xyxy[2] - xyxy[0]) / im0.shape[1]
        height = (xyxy[3] - xyxy[1]) / im0.shape[0]
        result = {
            'class': int(cls),
            'x_center': float(x_center),
            'y_center': float(y_center),
            'width': float(width),
            'height': float(height),
            'confidence': float(conf)
        }
        results.append(result)
    with open(json_path, 'w') as f:
        json.dump(results, f, indent=4)

3.2 運行推理腳本

在修改完 detect.py 后，我們可以通過以下命令運行推理腳本，并指定保存 json 文件的路徑：

python detect.py --source path/to/images --save-json --output path/to/output

其中，--source 指定輸入圖像或視頻的路徑，--save-json 表示保存 json 文件，--output 指定輸出路徑。

4. 總結

通過修改 YOLOv5 的推理腳本 detect.py，我們可以輕松地將預測結果保存為 txt 或 json 格式。txt 格式適合簡單的數據存儲，而 json 格式則更適合存儲結構化數據。根據實際需求選擇合適的輸出格式，可以大大提高后續數據處理的效率。

希望本文對你理解和使用 YOLOv5 有所幫助。如果你有任何問題或建議，歡迎在評論區留言討論。