Reference for ultralytics/engine/results.py

def __getitem__(self, idx):
    """
    Returns a new BaseTensor instance containing the specified indexed elements of the data tensor.

    Args:
        idx (int | List[int] | torch.Tensor): Index or indices to select from the data tensor.

    Returns:
        (BaseTensor): A new BaseTensor instance containing the indexed data.

    Examples:
        >>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
        >>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
        >>> result = base_tensor[0]  # Select the first row
        >>> print(result.data)
        tensor([1, 2, 3])
    """
    return self.__class__(self.data[idx], self.orig_shape)

def __len__(self):  # override len(results)
    """
    Returns the length of the underlying data tensor.

    Returns:
        (int): The number of elements in the first dimension of the data tensor.

    Examples:
        >>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
        >>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
        >>> len(base_tensor)
        2
    """
    return len(self.data)

def cpu(self):
    """
    Returns a copy of the tensor stored in CPU memory.

    Returns:
        (BaseTensor): A new BaseTensor object with the data tensor moved to CPU memory.

    Examples:
        >>> data = torch.tensor([[1, 2, 3], [4, 5, 6]]).cuda()
        >>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
        >>> cpu_tensor = base_tensor.cpu()
        >>> isinstance(cpu_tensor, BaseTensor)
        True
        >>> cpu_tensor.data.device
        device(type='cpu')
    """
    return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.cpu(), self.orig_shape)

def cuda(self):
    """
    Moves the tensor to GPU memory.

    Returns:
        (BaseTensor): A new BaseTensor instance with the data moved to GPU memory if it's not already a
            numpy array, otherwise returns self.

    Examples:
        >>> import torch
        >>> from ultralytics.engine.results import BaseTensor
        >>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
        >>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
        >>> gpu_tensor = base_tensor.cuda()
        >>> print(gpu_tensor.data.device)
        cuda:0
    """
    return self.__class__(torch.as_tensor(self.data).cuda(), self.orig_shape)

def numpy(self):
    """
    Returns a copy of the tensor as a numpy array.

    Returns:
        (np.ndarray): A numpy array containing the same data as the original tensor.

    Examples:
        >>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
        >>> orig_shape = (720, 1280)
        >>> base_tensor = BaseTensor(data, orig_shape)
        >>> numpy_array = base_tensor.numpy()
        >>> print(type(numpy_array))
        <class 'numpy.ndarray'>
    """
    return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.numpy(), self.orig_shape)

def to(self, *args, **kwargs):
    """
    Return a copy of the tensor with the specified device and dtype.

    Args:
        *args (Any): Variable length argument list to be passed to torch.Tensor.to().
        **kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().

    Returns:
        (BaseTensor): A new BaseTensor instance with the data moved to the specified device and/or dtype.

    Examples:
        >>> base_tensor = BaseTensor(torch.randn(3, 4), orig_shape=(480, 640))
        >>> cuda_tensor = base_tensor.to("cuda")
        >>> float16_tensor = base_tensor.to(dtype=torch.float16)
    """
    return self.__class__(torch.as_tensor(self.data).to(*args, **kwargs), self.orig_shape)

def __getitem__(self, idx):
    """
    Return a Results object for a specific index of inference results.

    Args:
        idx (int | slice): Index or slice to retrieve from the Results object.

    Returns:
        (Results): A new Results object containing the specified subset of inference results.

    Examples:
        >>> results = model("path/to/image.jpg")  # Perform inference
        >>> single_result = results[0]  # Get the first result
        >>> subset_results = results[1:4]  # Get a slice of results
    """
    return self._apply("__getitem__", idx)

def __len__(self):
    """
    Return the number of detections in the Results object.

    Returns:
        (int): The number of detections, determined by the length of the first non-empty attribute
            (boxes, masks, probs, keypoints, or obb).

    Examples:
        >>> results = Results(orig_img, path, names, boxes=torch.rand(5, 4))
        >>> len(results)
        5
    """
    for k in self._keys:
        v = getattr(self, k)
        if v is not None:
            return len(v)

def cpu(self):
    """
    Returns a copy of the Results object with all its tensors moved to CPU memory.

    This method creates a new Results object with all tensor attributes (boxes, masks, probs, keypoints, obb)
    transferred to CPU memory. It's useful for moving data from GPU to CPU for further processing or saving.

    Returns:
        (Results): A new Results object with all tensor attributes on CPU memory.

    Examples:
        >>> results = model("path/to/image.jpg")  # Perform inference
        >>> cpu_result = results[0].cpu()  # Move the first result to CPU
        >>> print(cpu_result.boxes.device)  # Output: cpu
    """
    return self._apply("cpu")

def cuda(self):
    """
    Moves all tensors in the Results object to GPU memory.

    Returns:
        (Results): A new Results object with all tensors moved to CUDA device.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> cuda_results = results[0].cuda()  # Move first result to GPU
        >>> for result in results:
        ...     result_cuda = result.cuda()  # Move each result to GPU
    """
    return self._apply("cuda")

def new(self):
    """
    Creates a new Results object with the same image, path, names, and speed attributes.

    Returns:
        (Results): A new Results object with copied attributes from the original instance.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> new_result = results[0].new()
    """
    return Results(orig_img=self.orig_img, path=self.path, names=self.names, speed=self.speed)

def numpy(self):
    """
    Converts all tensors in the Results object to numpy arrays.

    Returns:
        (Results): A new Results object with all tensors converted to numpy arrays.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> numpy_result = results[0].numpy()
        >>> type(numpy_result.boxes.data)
        <class 'numpy.ndarray'>

    Notes:
        This method creates a new Results object, leaving the original unchanged. It's useful for
        interoperability with numpy-based libraries or when CPU-based operations are required.
    """
    return self._apply("numpy")

def plot(
    self,
    conf=True,
    line_width=None,
    font_size=None,
    font="Arial.ttf",
    pil=False,
    img=None,
    im_gpu=None,
    kpt_radius=5,
    kpt_line=True,
    labels=True,
    boxes=True,
    masks=True,
    probs=True,
    show=False,
    save=False,
    filename=None,
    color_mode="class",
):
    """
    Plots detection results on an input RGB image.

    Args:
        conf (bool): Whether to plot detection confidence scores.
        line_width (float | None): Line width of bounding boxes. If None, scaled to image size.
        font_size (float | None): Font size for text. If None, scaled to image size.
        font (str): Font to use for text.
        pil (bool): Whether to return the image as a PIL Image.
        img (np.ndarray | None): Image to plot on. If None, uses original image.
        im_gpu (torch.Tensor | None): Normalized image on GPU for faster mask plotting.
        kpt_radius (int): Radius of drawn keypoints.
        kpt_line (bool): Whether to draw lines connecting keypoints.
        labels (bool): Whether to plot labels of bounding boxes.
        boxes (bool): Whether to plot bounding boxes.
        masks (bool): Whether to plot masks.
        probs (bool): Whether to plot classification probabilities.
        show (bool): Whether to display the annotated image.
        save (bool): Whether to save the annotated image.
        filename (str | None): Filename to save image if save is True.
        color_mode (bool): Specify the color mode, e.g., 'instance' or 'class'. Default to 'class'.

    Returns:
        (np.ndarray): Annotated image as a numpy array.

    Examples:
        >>> results = model("image.jpg")
        >>> for result in results:
        ...     im = result.plot()
        ...     im.show()
    """
    assert color_mode in {"instance", "class"}, f"Expected color_mode='instance' or 'class', not {color_mode}."
    if img is None and isinstance(self.orig_img, torch.Tensor):
        img = (self.orig_img[0].detach().permute(1, 2, 0).contiguous() * 255).to(torch.uint8).cpu().numpy()

    names = self.names
    is_obb = self.obb is not None
    pred_boxes, show_boxes = self.obb if is_obb else self.boxes, boxes
    pred_masks, show_masks = self.masks, masks
    pred_probs, show_probs = self.probs, probs
    annotator = Annotator(
        deepcopy(self.orig_img if img is None else img),
        line_width,
        font_size,
        font,
        pil or (pred_probs is not None and show_probs),  # Classify tasks default to pil=True
        example=names,
    )

    # Plot Segment results
    if pred_masks and show_masks:
        if im_gpu is None:
            img = LetterBox(pred_masks.shape[1:])(image=annotator.result())
            im_gpu = (
                torch.as_tensor(img, dtype=torch.float16, device=pred_masks.data.device)
                .permute(2, 0, 1)
                .flip(0)
                .contiguous()
                / 255
            )
        idx = (
            pred_boxes.id
            if pred_boxes.id is not None and color_mode == "instance"
            else pred_boxes.cls
            if pred_boxes and color_mode == "class"
            else reversed(range(len(pred_masks)))
        )
        annotator.masks(pred_masks.data, colors=[colors(x, True) for x in idx], im_gpu=im_gpu)

    # Plot Detect results
    if pred_boxes is not None and show_boxes:
        for i, d in enumerate(reversed(pred_boxes)):
            c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
            name = ("" if id is None else f"id:{id} ") + names[c]
            label = (f"{name} {conf:.2f}" if conf else name) if labels else None
            box = d.xyxyxyxy.reshape(-1, 4, 2).squeeze() if is_obb else d.xyxy.squeeze()
            annotator.box_label(
                box,
                label,
                color=colors(
                    c
                    if color_mode == "class"
                    else id
                    if id is not None
                    else i
                    if color_mode == "instance"
                    else None,
                    True,
                ),
                rotated=is_obb,
            )

    # Plot Classify results
    if pred_probs is not None and show_probs:
        text = ",\n".join(f"{names[j] if names else j} {pred_probs.data[j]:.2f}" for j in pred_probs.top5)
        x = round(self.orig_shape[0] * 0.03)
        annotator.text([x, x], text, txt_color=(255, 255, 255))  # TODO: allow setting colors

    # Plot Pose results
    if self.keypoints is not None:
        for i, k in enumerate(reversed(self.keypoints.data)):
            annotator.kpts(
                k,
                self.orig_shape,
                radius=kpt_radius,
                kpt_line=kpt_line,
                kpt_color=colors(i, True) if color_mode == "instance" else None,
            )

    # Show results
    if show:
        annotator.show(self.path)

    # Save results
    if save:
        annotator.save(filename)

    return annotator.result()

def save(self, filename=None, *args, **kwargs):
    """
    Saves annotated inference results image to file.

    This method plots the detection results on the original image and saves the annotated image to a file. It
    utilizes the `plot` method to generate the annotated image and then saves it to the specified filename.

    Args:
        filename (str | Path | None): The filename to save the annotated image. If None, a default filename
            is generated based on the original image path.
        *args (Any): Variable length argument list to be passed to the `plot` method.
        **kwargs (Any): Arbitrary keyword arguments to be passed to the `plot` method.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> for result in results:
        ...     result.save("annotated_image.jpg")
        >>> # Or with custom plot arguments
        >>> for result in results:
        ...     result.save("annotated_image.jpg", conf=False, line_width=2)
    """
    if not filename:
        filename = f"results_{Path(self.path).name}"
    self.plot(save=True, filename=filename, *args, **kwargs)
    return filename

def save_crop(self, save_dir, file_name=Path("im.jpg")):
    """
    Saves cropped detection images to specified directory.

    This method saves cropped images of detected objects to a specified directory. Each crop is saved in a
    subdirectory named after the object's class, with the filename based on the input file_name.

    Args:
        save_dir (str | Path): Directory path where cropped images will be saved.
        file_name (str | Path): Base filename for the saved cropped images. Default is Path("im.jpg").

    Notes:
        - This method does not support Classify or Oriented Bounding Box (OBB) tasks.
        - Crops are saved as 'save_dir/class_name/file_name.jpg'.
        - The method will create necessary subdirectories if they don't exist.
        - Original image is copied before cropping to avoid modifying the original.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> for result in results:
        ...     result.save_crop(save_dir="path/to/crops", file_name="detection")
    """
    if self.probs is not None:
        LOGGER.warning("WARNING ⚠️ Classify task do not support `save_crop`.")
        return
    if self.obb is not None:
        LOGGER.warning("WARNING ⚠️ OBB task do not support `save_crop`.")
        return
    for d in self.boxes:
        save_one_box(
            d.xyxy,
            self.orig_img.copy(),
            file=Path(save_dir) / self.names[int(d.cls)] / f"{Path(file_name)}.jpg",
            BGR=True,
        )

def save_txt(self, txt_file, save_conf=False):
    """
    Save detection results to a text file.

    Args:
        txt_file (str | Path): Path to the output text file.
        save_conf (bool): Whether to include confidence scores in the output.

    Returns:
        (str): Path to the saved text file.

    Examples:
        >>> from ultralytics import YOLO
        >>> model = YOLO("yolo11n.pt")
        >>> results = model("path/to/image.jpg")
        >>> for result in results:
        ...     result.save_txt("output.txt")

    Notes:
        - The file will contain one line per detection or classification with the following structure:
          - For detections: `class confidence x_center y_center width height`
          - For classifications: `confidence class_name`
          - For masks and keypoints, the specific formats will vary accordingly.
        - The function will create the output directory if it does not exist.
        - If save_conf is False, the confidence scores will be excluded from the output.
        - Existing contents of the file will not be overwritten; new results will be appended.
    """
    is_obb = self.obb is not None
    boxes = self.obb if is_obb else self.boxes
    masks = self.masks
    probs = self.probs
    kpts = self.keypoints
    texts = []
    if probs is not None:
        # Classify
        [texts.append(f"{probs.data[j]:.2f} {self.names[j]}") for j in probs.top5]
    elif boxes:
        # Detect/segment/pose
        for j, d in enumerate(boxes):
            c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
            line = (c, *(d.xyxyxyxyn.view(-1) if is_obb else d.xywhn.view(-1)))
            if masks:
                seg = masks[j].xyn[0].copy().reshape(-1)  # reversed mask.xyn, (n,2) to (n*2)
                line = (c, *seg)
            if kpts is not None:
                kpt = torch.cat((kpts[j].xyn, kpts[j].conf[..., None]), 2) if kpts[j].has_visible else kpts[j].xyn
                line += (*kpt.reshape(-1).tolist(),)
            line += (conf,) * save_conf + (() if id is None else (id,))
            texts.append(("%g " * len(line)).rstrip() % line)

    if texts:
        Path(txt_file).parent.mkdir(parents=True, exist_ok=True)  # make directory
        with open(txt_file, "a") as f:
            f.writelines(text + "\n" for text in texts)

def show(self, *args, **kwargs):
    """
    Display the image with annotated inference results.

    This method plots the detection results on the original image and displays it. It's a convenient way to
    visualize the model's predictions directly.

    Args:
        *args (Any): Variable length argument list to be passed to the `plot()` method.
        **kwargs (Any): Arbitrary keyword arguments to be passed to the `plot()` method.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> results[0].show()  # Display the first result
        >>> for result in results:
        ...     result.show()  # Display all results
    """
    self.plot(show=True, *args, **kwargs)

def summary(self, normalize=False, decimals=5):
    """
    Converts inference results to a summarized dictionary with optional normalization for box coordinates.

    This method creates a list of detection dictionaries, each containing information about a single
    detection or classification result. For classification tasks, it returns the top class and its
    confidence. For detection tasks, it includes class information, bounding box coordinates, and
    optionally mask segments and keypoints.

    Args:
        normalize (bool): Whether to normalize bounding box coordinates by image dimensions. Defaults to False.
        decimals (int): Number of decimal places to round the output values to. Defaults to 5.

    Returns:
        (List[Dict]): A list of dictionaries, each containing summarized information for a single
            detection or classification result. The structure of each dictionary varies based on the
            task type (classification or detection) and available information (boxes, masks, keypoints).

    Examples:
        >>> results = model("image.jpg")
        >>> summary = results[0].summary()
        >>> print(summary)
    """
    # Create list of detection dictionaries
    results = []
    if self.probs is not None:
        class_id = self.probs.top1
        results.append(
            {
                "name": self.names[class_id],
                "class": class_id,
                "confidence": round(self.probs.top1conf.item(), decimals),
            }
        )
        return results

    is_obb = self.obb is not None
    data = self.obb if is_obb else self.boxes
    h, w = self.orig_shape if normalize else (1, 1)
    for i, row in enumerate(data):  # xyxy, track_id if tracking, conf, class_id
        class_id, conf = int(row.cls), round(row.conf.item(), decimals)
        box = (row.xyxyxyxy if is_obb else row.xyxy).squeeze().reshape(-1, 2).tolist()
        xy = {}
        for j, b in enumerate(box):
            xy[f"x{j + 1}"] = round(b[0] / w, decimals)
            xy[f"y{j + 1}"] = round(b[1] / h, decimals)
        result = {"name": self.names[class_id], "class": class_id, "confidence": conf, "box": xy}
        if data.is_track:
            result["track_id"] = int(row.id.item())  # track ID
        if self.masks:
            result["segments"] = {
                "x": (self.masks.xy[i][:, 0] / w).round(decimals).tolist(),
                "y": (self.masks.xy[i][:, 1] / h).round(decimals).tolist(),
            }
        if self.keypoints is not None:
            x, y, visible = self.keypoints[i].data[0].cpu().unbind(dim=1)  # torch Tensor
            result["keypoints"] = {
                "x": (x / w).numpy().round(decimals).tolist(),  # decimals named argument required
                "y": (y / h).numpy().round(decimals).tolist(),
                "visible": visible.numpy().round(decimals).tolist(),
            }
        results.append(result)

    return results

def to(self, *args, **kwargs):
    """
    Moves all tensors in the Results object to the specified device and dtype.

    Args:
        *args (Any): Variable length argument list to be passed to torch.Tensor.to().
        **kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().

    Returns:
        (Results): A new Results object with all tensors moved to the specified device and dtype.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> result_cuda = results[0].to("cuda")  # Move first result to GPU
        >>> result_cpu = results[0].to("cpu")  # Move first result to CPU
        >>> result_half = results[0].to(dtype=torch.float16)  # Convert first result to half precision
    """
    return self._apply("to", *args, **kwargs)

def to_csv(self, normalize=False, decimals=5, *args, **kwargs):
    """
    Converts detection results to a CSV format.

    This method serializes the detection results into a CSV format. It includes information
    about detected objects such as bounding boxes, class names, confidence scores, and optionally
    segmentation masks and keypoints.

    Args:
        normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
            If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
        decimals (int): Number of decimal places to round the output values to. Defaults to 5.
        *args (Any): Variable length argument list to be passed to pandas.DataFrame.to_csv().
        **kwargs (Any): Arbitrary keyword arguments to be passed to pandas.DataFrame.to_csv().


    Returns:
        (str): CSV containing all the information in results in an organized way.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> csv_result = results[0].to_csv()
        >>> print(csv_result)
    """
    return self.to_df(normalize=normalize, decimals=decimals).to_csv(*args, **kwargs)

def to_df(self, normalize=False, decimals=5):
    """
    Converts detection results to a Pandas Dataframe.

    This method converts the detection results into Pandas Dataframe format. It includes information
    about detected objects such as bounding boxes, class names, confidence scores, and optionally
    segmentation masks and keypoints.

    Args:
        normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
            If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
        decimals (int): Number of decimal places to round the output values to. Defaults to 5.

    Returns:
        (DataFrame): A Pandas Dataframe containing all the information in results in an organized way.

    Examples:
        >>> results = model("path/to/image.jpg")
        >>> df_result = results[0].to_df()
        >>> print(df_result)
    """
    import pandas as pd

    return pd.DataFrame(self.summary(normalize=normalize, decimals=decimals))