问题
在使用super_gradients库中的Yolo-nas预测图片时，想要获取预测好的图片，但执行out = model.predict("camera01.png", conf=0.6, batch_size=None)之后只能out.show()和out.save()，无法返回预测结果图片。
解决方法
import cv2
import torch
from super_gradients.training import models
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
# yolo_nas_m,yolo_nas_l
model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)
out = model.predict("camera01.png", conf=0.6)
# 关键在这
sigle_image_out = out._images_prediction_lst[0]
image = sigle_image_out.draw()
cv2.imshow('image',image)
cv2.waitKey()原理
在out = model.predict("camera01.png", conf=0.6)处打断点， debug，进入到predict函数里:
def predict(
self,
images: ImageSource,
iou: Optional[float] = None,
conf: Optional[float] = None,
batch_size: int = 32,
fuse_model: bool = True,
) -> ImagesDetectionPrediction:
"""Predict an image or a list of images.
:param images: Images to predict.
:param iou: (Optional) IoU threshold for the nms algorithm. If None, the default value associated to the training is used.
:param conf: (Optional) Below the confidence threshold, prediction are discarded.
If None, the default value associated to the training is used.
:param batch_size: Maximum number of images to process at the same time.
:param fuse_model: If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
"""
pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model)
return pipeline(images, batch_size=batch_size) # type: ignore接着debug有函数的地方就进去，到最后发现，out是ImagesDetectionPrediction类
ImagesDetectionPrediction类的属性是_images_prediction_lst: List[ImageDetectionPrediction]
里面有多个ImageDetectionPrediction类。注意ImagesDetectionPrediction和ImageDetectionPrediction，一个是处理多个图片，一个是单个图片。其中ImageDetectionPrediction类有个draw方法，该方法画出了预测框并返回画出了预测框的结果。
所以如下两句代码，成功获取到了预测结果图片。sigle_image_out = out._images_prediction_lst[0] image = sigle_image_out.draw()
ImagesDetectionPrediction和ImageDetectionPrediction源码如下:
@dataclass
class ImagesDetectionPrediction(ImagesPredictions):
"""Object wrapping the list of image detection predictions.
:attr _images_prediction_lst: List of the predictions results
"""
_images_prediction_lst: List[ImageDetectionPrediction]
def show(self, box_thickness: int = 2, show_confidence: bool = True, color_mapping: Optional[List[Tuple[int, int, int]]] = None) -> None:
"""Display the predicted bboxes on the images.
:param box_thickness: Thickness of bounding boxes.
:param show_confidence: Whether to show confidence scores on the image.
:param color_mapping: List of tuples representing the colors for each class.
Default is None, which generates a default color mapping based on the number of class names.
"""
for prediction in self._images_prediction_lst:
prediction.show(box_thickness=box_thickness, show_confidence=show_confidence, color_mapping=color_mapping)
def save(
self, output_folder: str, box_thickness: int = 2, show_confidence: bool = True, color_mapping: Optional[List[Tuple[int, int, int]]] = None
) -> None:
"""Save the predicted bboxes on the images.
:param output_folder: Folder path, where the images will be saved.
:param box_thickness: Thickness of bounding boxes.
:param show_confidence: Whether to show confidence scores on the image.
:param color_mapping: List of tuples representing the colors for each class.
Default is None, which generates a default color mapping based on the number of class names.
"""
if output_folder:
os.makedirs(output_folder, exist_ok=True)
for i, prediction in enumerate(self._images_prediction_lst):
image_output_path = os.path.join(output_folder, f"pred_{i}.jpg")
prediction.save(output_path=image_output_path, box_thickness=box_thickness, show_confidence=show_confidence, color_mapping=color_mapping)_images_prediction_lst: List[ImageDetectionPrediction]
@dataclass
class ImageDetectionPrediction(ImagePrediction):
"""Object wrapping an image and a detection model's prediction.
:attr image: Input image
:attr predictions: Predictions of the model
:attr class_names: List of the class names to predict
"""
image: np.ndarray
prediction: DetectionPrediction
class_names: List[str]
def draw(self, box_thickness: int = 2, show_confidence: bool = True, color_mapping: Optional[List[Tuple[int, int, int]]] = None) -> np.ndarray:
"""Draw the predicted bboxes on the image.
:param box_thickness: Thickness of bounding boxes.
:param show_confidence: Whether to show confidence scores on the image.
:param color_mapping: List of tuples representing the colors for each class.
Default is None, which generates a default color mapping based on the number of class names.
:return: Image with predicted bboxes. Note that this does not modify the original image.
"""
image = self.image.copy()
color_mapping = color_mapping or generate_color_mapping(len(self.class_names))
for pred_i in np.argsort(self.prediction.confidence):
class_id = int(self.prediction.labels[pred_i])
score = "" if not show_confidence else str(round(self.prediction.confidence[pred_i], 2))
image = draw_bbox(
image=image,
title=f"{self.class_names[class_id]} {score}",
color=color_mapping[class_id],
box_thickness=box_thickness,
x1=int(self.prediction.bboxes_xyxy[pred_i, 0]),
y1=int(self.prediction.bboxes_xyxy[pred_i, 1]),
x2=int(self.prediction.bboxes_xyxy[pred_i, 2]),
y2=int(self.prediction.bboxes_xyxy[pred_i, 3]),
)
return image
def show(self, box_thickness: int = 2, show_confidence: bool = True, color_mapping: Optional[List[Tuple[int, int, int]]] = None) -> None:
"""Display the image with predicted bboxes.
:param box_thickness: Thickness of bounding boxes.
:param show_confidence: Whether to show confidence scores on the image.
:param color_mapping: List of tuples representing the colors for each class.
Default is None, which generates a default color mapping based on the number of class names.
"""
image = self.draw(box_thickness=box_thickness, show_confidence=show_confidence, color_mapping=color_mapping)
show_image(image)
def save(self, output_path: str, box_thickness: int = 2, show_confidence: bool = True, color_mapping: Optional[List[Tuple[int, int, int]]] = None) -> None:
"""Save the predicted bboxes on the images.
:param output_path: Path to the output video file.
:param box_thickness: Thickness of bounding boxes.
:param show_confidence: Whether to show confidence scores on the image.
:param color_mapping: List of tuples representing the colors for each class.
Default is None, which generates a default color mapping based on the number of class names.
"""
image = self.draw(box_thickness=box_thickness, show_confidence=show_confidence, color_mapping=color_mapping)
save_image(image=image, path=output_path)