Experimenting Yolov5 segmentation evaluation with prediction saved txt files as input

[yuewang199511]

Hi, I am running an experience to save the predictions as txt files and then use (almost) the same logic in evaluation task to compare the ground truth with them.

I have made the revised evaluation code running successfully the ground truth vs ground truth seems fine, but found some problems that make me wonder.
I found there are two problems:

1. evaluation module has no working method to save txt format for segmentation while prediction has it. Both in yolov5 and ultralytics. Why?
2. ground truth vs prediction metrics are lower than the it shows in the yolov5 segment/val.py module

The revised code is attached below and here are the commands you can reproduce the process:

get prediction labels saved
python3 segment/predict.py --weights yolov5s-seg.pt --data coco128-seg.yaml --save-txt --save-conf

get evaluation result
'python segment/val.py --weights yolov5s-seg.pt --data coco128-seg.yaml --verbose'

get the evaluation result from labels and predictions
python3 evaluation_seg.py {annotated dataset dir} {predicted labels dir} {dir to save evaluation graphs later}

import os
import torch
import numpy as np
from pathlib import Path
from tqdm import tqdm
import json
import torchvision.transforms as transforms
import argparse

import torch.nn.functional as F
from PIL import Image, ImageOps
from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
from utils.general import (
    LOGGER,
    TQDM_BAR_FORMAT,
    Profile,
    check_dataset,
    check_img_size,
    check_requirements,
    check_yaml,
    coco80_to_coco91_class,
    colorstr,
    increment_path,
    non_max_suppression,
    print_args,
    scale_boxes,
    xywh2xyxy,
    xyxy2xywh,
    xywhn2xyxy,
    xyn2xy,
    segments2boxes
)
from utils.plots import output_to_target, plot_images, plot_val_study
from utils.segment.dataloaders import create_dataloader, polygons2masks
from utils.segment.general import mask_iou, process_mask, process_mask_native, scale_image
from utils.segment.metrics import Metrics, ap_per_class_box_and_mask
from utils.segment.plots import plot_images_and_masks
import yaml
import cv2
from reportlab.pdfgen import canvas



# 将所有图片合并为一个pdf
def generate_report_file(model_dir):

    # Get a list of all PNG and JPG files in the folder
    image_files = [f for f in os.listdir(model_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    # Sort files to maintain a specific order (optional)
    image_files.sort()

    # Create a PDF file
    pdf_path = os.path.join(model_dir, 'report.pdf')
    c = canvas.Canvas(pdf_path)

    for image_file in image_files:
        image_path = os.path.join(model_dir, image_file)
        img = Image.open(image_path)
        
        # Calculate the aspect ratio to fit the image within the page
        width, height = img.size
        aspect = height / float(width)
        
        # Set the size of the page to the size of the image
        page_width = 595.27  # A4 page width in points
        page_height = page_width * aspect
        
        # Resize the canvas to match the image size
        c.setPageSize((page_width, page_height))
        
        # Draw the image on the canvas
        c.drawImage(image_path, 0, 0, width=page_width, height=page_height)
        c.showPage()  # Add a new page

    # Save the PDF
    c.save()
    print(f"PDF created successfully: {pdf_path}")
    return pdf_path


def reverse_scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
    """Scales bounding boxes from img0_shape to img1_shape, optionally using provided `ratio_pad`."""
    if ratio_pad is None:  # calculate from img0_shape
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain = old / new
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]

    boxes[..., :4] *= gain  # Scale the boxes by gain
    boxes[..., [0, 2]] += pad[0]  # Add the x padding
    boxes[..., [1, 3]] += pad[1]  # Add the y padding
    return boxes

def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
    """
    Return correct prediction matrix
    Arguments:
        detections (array[N, 6]), x1, y1, x2, y2, conf, class
        labels (array[M, 5]), class, x1, y1, x2, y2
    Returns:
        correct (array[N, 10]), for 10 IoU levels.
    """
    if masks:
        if overlap:
            nl = len(labels)
            index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
            gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
            gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
        if gt_masks.shape[1:] != pred_masks.shape[1:]:
            gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
            gt_masks = gt_masks.gt_(0.5)
        iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
    else:  # boxes
        iou = box_iou(labels[:, 1:], detections[:, :4])
    correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
    correct_class = labels[:, 0:1] == detections[:, 5]
    for i in range(len(iouv)):
        x = torch.where((iou >= iouv[i]) & correct_class)  # IoU > threshold and classes match
        if x[0].shape[0]:
            matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detect, iou]
            if x[0].shape[0] > 1:
                matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                # matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
            correct[matches[:, 1].astype(int), i] = True
    return torch.tensor(correct, dtype=torch.bool, device=iouv.device)


def generate_report(annotation_dir, prediction_dir, save_dir):
    # Set paths
    overlap = False
    downsample_ratio =1
    os.makedirs(save_dir, exist_ok=True)

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    verbose = True
    iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for [email protected]:0.95
    niou = iouv.numel()
    # Load class names
    if os.path.exists(annotation_dir/ 'classes.txt'):
        with open(annotation_dir/ 'classes.txt', 'r') as f:
            class_names = f.read().splitlines()
    elif len(list(annotation_dir.glob('*.yaml'))) == 1:
        with open(list(annotation_dir.glob('*.yaml'))[0]) as f:
            data = yaml.safe_load(f)
            class_names = data['names']
    else:
        raise Exception('You should have a classes.txt file or any .yaml file with class names in {}!'.format(annotation_dir))
    nc = len(class_names)  # number of classes

    # Initialize metrics
    seen, plots = 0, True  # number of IoU thresholds for metrics
    jdict, stats, ap, ap_class = [], [], [], []
    confusion_matrix = ConfusionMatrix(nc=nc)
    imgsz = 640
    batch_size = 32
    stride = 32
    single_cls = True if nc == 1 else False
    pad, rect = (0.5, True)
    metrics = Metrics()

    dataloader = create_dataloader(
            annotation_dir,
            imgsz,
            batch_size,
            stride,
            single_cls,
            pad=pad,
            rect=rect,
            workers=8,
            prefix=colorstr(f"val: "),
            overlap_mask=overlap,
            mask_downsample_ratio=1,
        )[0]
    s = ("%22s" + "%11s" * 10) % (
        "Class",
        "Images",
        "Instances",
        "Box(P",
        "R",
        "mAP50",
        "mAP50-95)",
        "Mask(P",
        "R",
        "mAP50",
        "mAP50-95)",
    )
    tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    dt = Profile(device=device), Profile(device=device), Profile(device=device)  # profiling times
    cuda = device.type != "cpu"
    pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT)  # progress bar
    for batch_i, (im, targets, paths, shapes, masks) in enumerate(pbar):
        # Get corresponding label and prediction file paths
        with dt[0]:
            if cuda:
                im = im.to(device, non_blocking=True)
                targets = targets.to(device)
                masks = masks.to(device)
            masks = masks.float()
            im = im.half()
            im /= 255  # 0 - 255 to 0.0 - 1.0
            nb, _, height, width = im.shape  # batch size, channels, height, width
        targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)  # to pixels
        preds = []
        segments = []

        for path in paths:
            pred_path = prediction_dir / f"{Path(path).stem}.txt"
            if not os.path.exists(pred_path):
                pred = torch.empty((0, 6), dtype=torch.float32, device=device)
                segment = []
            else:
                pred = []
                segment = []
                with open(pred_path, 'r') as f:
                    for line in f.readlines():
                        # Split the line into columns
                        line = [float(element) for element in line.split()]
                        if len(line) <= 2:
                            continue
                        if len(line[1:]) % 2 != 0:
                            seg = np.array(line[1:-1]).reshape(-1,2)
                            pred.append([line[0], *segments2boxes([seg])[0], line[-1]])
                            segment.append(seg)
                        else:
                            seg = np.array(line[1:]).reshape(-1, 2)
                            pred.append([line[0], *segments2boxes([seg])[0], 1.0])
                            segment.append(seg)
                # Reshape the result (if needed)
                pred = np.array(pred)
                pred = pred[:, [1, 2, 3, 4, 5, 0]]  # Rearrange columns
                pred = torch.tensor(pred, device=device)
            preds.append(pred)
            segments.append(segment)
        plot_masks = []  # masks for plotting
        # Metrics
        for si, (pred, segment) in enumerate(zip(preds, segments)):
            labels = targets[targets[:, 0] == si, 1:]
            nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
            path, shape = Path(paths[si]), shapes[si][0]
            correct_masks = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
            correct_bboxes = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
            seen += 1

            if npr == 0:
                if nl:
                    stats.append((correct_masks, correct_bboxes, *torch.zeros((2, 0), device=device), labels[:, 0]))
                    if plots:
                        confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
                continue
            # Masks
            for seg in segment:
                seg[:, 0] *= shape[1]*shapes[si][1][0][0]
                seg[:, 1] *= shape[0]*shapes[si][1][0][0]
                seg[:, 0] += shapes[si][1][1][0]
                seg[:, 1] += shapes[si][1][1][1]
            pred_masks = polygons2masks(im[si].shape[1:], segment, 1, downsample_ratio=1)
            pred_masks = (
            torch.as_tensor(pred_masks, device=device, dtype=torch.float32)
            if len(pred_masks)
            else torch.zeros(
                1 if overlap else nl, im[si].shape[1] // downsample_ratio, im[si].shape[2] // downsample_ratio, device=device, dtype=torch.float32
            )
        )
            if len(pred_masks) != npr:
                additional_masks = torch.zeros(
                    npr - len(pred_masks), im[si].shape[1] // downsample_ratio, im[si].shape[2] // downsample_ratio,
                    device=device, dtype=torch.float32
                )
                pred_masks = torch.cat([pred_masks, additional_masks], dim=0)
            # Predictions
            if single_cls:
                pred[:, 5] = 0
            midx = [si] if overlap else targets[:, 0] == si
            gt_masks = masks[midx]
            predn = pred.clone()
            predn = xywhn2xyxy(predn, w=shapes[si][0][1], h=shapes[si][0][0])
            # Evaluate
            if nl:
                tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
                scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
                labelsn = torch.cat((labels[:, 0:1], tbox), 1)  # native-space labels
                correct_bboxes = process_batch(predn, labelsn, iouv)
                correct_masks = process_batch(predn, labelsn, iouv, pred_masks, gt_masks, overlap=overlap, masks=True)
                if plots:
                    confusion_matrix.process_batch(predn, labelsn)
            stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:, 5], labels[:, 0]))  # (conf, pcls, tcls)
            # To inference image shape
            pred = xywhn2xyxy(pred, w=shapes[si][0][1], h=shapes[si][0][0])
            reverse_scale_boxes(im[si].shape[1:], pred, shape, shapes[si][1])
            preds[si] = pred
            pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8)
            if plots and batch_i < 3:
                plot_masks.append(pred_masks[:15])  # filter top 15 to plot
        # Plot images
        if plots and batch_i < 3:
            if len(plot_masks):
                plot_masks = torch.cat(plot_masks, dim=0)
            plot_images_and_masks(im, targets, masks, paths, save_dir / f"val_batch{batch_i}_labels.jpg", class_names)
            plot_images_and_masks(
                im,
                output_to_target(preds, max_det=15),
                plot_masks,
                paths,
                save_dir / f"val_batch{batch_i}_pred.jpg",
                class_names,
            )  # pred


    # Compute metrics
    stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
    if len(stats) and stats[0].any():
        results = ap_per_class_box_and_mask(*stats, plot=plots, save_dir=save_dir, names=class_names)
        metrics.update(results)
    nt = np.bincount(stats[4].astype(int), minlength=nc)  # number of targets per class

    # Print results
    pf = "%22s" + "%11i" * 2 + "%11.3g" * 8  # print format
    LOGGER.info(pf % ("all", seen, nt.sum(), *metrics.mean_results()))
    if nt.sum() == 0:
        LOGGER.warning(f"WARNING ⚠️ no labels found in {annotation_dir} set, can not compute metrics without labels")

    metrics_dict = {}
    # Print results per class
    if len(stats):
        for i, c in enumerate(metrics.ap_class_index):
            result = metrics.class_result(i)
            LOGGER.info(pf % (class_names[c], seen, nt[c], *result))
            suffix = ('box', 'seg')
            for j in range(2):
                metrics_dict[f'precision/{suffix[j]}@{class_names[c]}'] = result[j * 4 + 0]
                metrics_dict [f'recall/{suffix[j]}@{class_names[c]}'] = result[j * 4 + 0]
                metrics_dict [f'ap50/{suffix[j]}@{class_names[c]}'] = result[j * 4 + 0]
                metrics_dict [f'ap/{suffix[j]}@{class_names[c]}'] = result[j * 4 + 0]

    # Plots
    if plots:
        confusion_matrix.plot(save_dir=save_dir, names=list(class_names.values()))
    mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask = metrics.mean_results()
    # Return results
    metrics_dict ['meam_precision_box'] = mp_bbox
    metrics_dict ['mean_recall_box'] = mr_bbox
    metrics_dict ['map50_box'] = map50_bbox
    metrics_dict ['map_box'] = map_bbox
    metrics_dict ['meam_precision_seg'] = mp_mask
    metrics_dict ['mean_recall_seg'] = mr_mask
    metrics_dict ['map50_seg'] = map50_mask
    metrics_dict ['map_seg'] = map_mask
    with open(save_dir/ 'metrics.json', 'w') as json_file:
        json.dump(metrics_dict, json_file)

    generate_report_file(save_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("annotation_dir", help="the directory for ground truth")
    parser.add_argument("prediction_dir", help="the directory for prediction results saved as txt files")
    parser.add_argument("save_dir", help='the output directory to save report files')

    opt = parser.parse_args()
    annotation_dir = Path(opt.annotation_dir)
    prediction_dir = Path(opt.prediction_dir)
    save_dir = Path(opt.save_dir)

    generate_report(annotation_dir, prediction_dir, save_dir)