### 将Matplotlib生成的图像转换为YOLO模型可预测的格式 为了使由 `matplotlib` 生成的图像能够被 YOLO 模型用于预测，需遵循特定的数据预处理步骤。这包括读取图像、调整大小以及标准化等操作[^1]。 #### 图像准备过程 首先，确保安装必要的库： ```bash pip install opencv-python numpy matplotlib torch torchvision ``` 接着，在Python脚本中实现如下代码片段完成图像格式转换： ```python import cv2 import numpy as np from PIL import Image import matplotlib.pyplot as plt import torch from torchvision.transforms.functional import to_tensor def plot_to_image(fig): """Converts a Matplotlib figure into an image.""" fig.canvas.draw() img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,)) return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) def preprocess_for_yolo(image_path=None, figsize=(6, 4), input_size=640): """ Converts and prepares the given or plotted image for YOLO prediction. :param image_path: Path of the source image file. If None, uses current plot. :param figsize: Size tuple for creating new plots with specified dimensions. :param input_size: Desired size after resizing images before feeding them into YOLO network. :return: Preprocessed tensor ready for passing through YOLO model. """ if not image_path: # Capture existing plot directly without saving it temporarily on disk fig = plt.gcf() if not fig.axes: raise ValueError("No active figure found.") img_bgr = plot_to_image(fig) else: # Load image using OpenCV when provided path instead of capturing screen content img_bgr = cv2.imread(image_path) # Resize while maintaining aspect ratio h, w = img_bgr.shape[:2] scale_factor = min(input_size / float(h), input_size / float(w)) resized_h = int(round(scale_factor * h)) resized_w = int(round(scale_factor * w)) padded_img = np.full((input_size, input_size, 3), fill_value=114, dtype=np.uint8) resized_img = cv2.resize(img_bgr, dsize=(resized_w, resized_h), interpolation=cv2.INTER_LINEAR) dh = (input_size - resized_h) // 2 dw = (input_size - resized_w) // 2 padded_img[dh:(dh+resized_h), dw:(dw+resized_w)] = resized_img # Convert HWC -> CHW & Normalize [0., 1.] then convert to Tensor normalized_img = ((padded_img.astype(np.float32)) / 255.).transpose(2, 0, 1)[np.newaxis,:,:,:] return torch.tensor(normalized_img).to('cuda' if torch.cuda.is_available() else 'cpu') # Example usage: plt.plot([1, 2], [3, 4]) # Create some simple line chart just for demonstration purposes here. preprocessed_data = preprocess_for_yolo(figsize=(6, 4)) print(preprocessed_data.size()) ``` 上述代码实现了两个主要函数：一个是将当前绘图窗口的内容转化为numpy数组形式；另一个则是针对给定路径下的图片或是直接获取屏幕上的图表并按照YOLO的要求做相应的前处理工作，最终返回一个可以送入GPU/CPU设备上运行的张量对象[^2]。