object-segmentation/tests/show_yolo_seg.py

#!/usr/bin/env python3
"""
show_yolo_seg.py

Usage:
    python show_yolo_seg.py /path/to/image.jpg /path/to/labels.txt

Supports:
 - Segmentation polygons: "class x1 y1 x2 y2 ... xn yn"
 - YOLO bbox lines as fallback: "class x_center y_center width height"
Coordinates can be normalized [0..1] or absolute pixels (auto-detected).
"""
import sys
import cv2
import numpy as np
import matplotlib.pyplot as plt
import argparse
from pathlib import Path
import random
from shapely.geometry import LineString

from src.utils.image import Image


def parse_label_line(line):
    parts = line.strip().split()
    if not parts:
        return None
    cls = int(float(parts[0]))
    coords = [float(x) for x in parts[1:]]
    return cls, coords


def coords_are_normalized(coords):
    # If every coordinate is between 0 and 1 (inclusive-ish), assume normalized
    if not coords:
        return False
    return max(coords) <= 1.001


def yolo_bbox_to_xyxy(coords, img_w, img_h):
    # coords: [xc, yc, w, h] normalized or absolute
    xc, yc, w, h = coords[:4]
    if max(coords) <= 1.001:
        xc *= img_w
        yc *= img_h
        w *= img_w
        h *= img_h
    x1 = int(round(xc - w / 2))
    y1 = int(round(yc - h / 2))
    x2 = int(round(xc + w / 2))
    y2 = int(round(yc + h / 2))
    return x1, y1, x2, y2


def poly_to_pts(coords, img_w, img_h):
    # coords: [x1 y1 x2 y2 ...] either normalized or absolute
    if coords_are_normalized(coords[4:]):
        coords = [coords[i] * (img_w if i % 2 == 0 else img_h) for i in range(len(coords))]
    pts = np.array(coords, dtype=np.int32).reshape(-1, 2)
    return pts


def random_color_for_class(cls):
    random.seed(cls)  # deterministic per class
    return (
        0,
        0,
        255,
    )  # tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))


def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
    # img: BGR numpy array
    overlay = img.copy()
    h, w = img.shape[:2]
    for line in labels:
        if isinstance(line, str):
            cls, coords = parse_label_line(line)
        if isinstance(line, tuple):
            cls, coords = line

        if not coords:
            continue
        # polygon case (>=6 coordinates)
        if len(coords) >= 6:
            color = random_color_for_class(cls)

            x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords[:4], w, h)
            print(x1, y1, x2, y2)
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)

            pts = poly_to_pts(coords[4:], w, h)
            # line = LineString(pts)
            # # Buffer distance in pixels
            # buffered = line.buffer(3, cap_style=2, join_style=2)
            # coords = np.array(buffered.exterior.coords, dtype=np.int32)
            # cv2.fillPoly(overlay, [coords], color=(255, 255, 255))

            # fill on overlay
            cv2.fillPoly(overlay, [pts], color)
            # outline on base image
            cv2.polylines(img, [pts], isClosed=True, color=color, thickness=1)
            # put class text at first point
            x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
            if 0:
                cv2.putText(
                    img,
                    str(cls),
                    (x, max(6, y)),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6,
                    (255, 255, 255),
                    2,
                    cv2.LINE_AA,
                )

        # YOLO bbox case (4 coords)
        elif len(coords) == 4:
            x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords, w, h)
            color = random_color_for_class(cls)
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv2.putText(
                img,
                str(cls),
                (x1, max(6, y1 - 4)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                2,
                cv2.LINE_AA,
            )
        else:
            # Unknown / invalid format, skip
            continue

    # blend overlay for filled polygons
    cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img)
    return img


def load_labels_file(label_path):
    labels = []
    with open(label_path, "r") as f:
        for raw in f:
            line = raw.strip()
            if not line:
                continue
            parsed = parse_label_line(line)
            if parsed:
                labels.append(parsed)
    return labels


def main():
    parser = argparse.ArgumentParser(description="Show YOLO segmentation / polygon annotations")
    parser.add_argument("image", type=str, help="Path to image file")
    parser.add_argument("--labels", type=str, help="Path to YOLO label file (polygons)")
    parser.add_argument("--alpha", type=float, default=0.4, help="Polygon fill alpha (0..1)")
    parser.add_argument("--no-bbox", action="store_true", help="Don't draw bounding boxes for polygons")
    args = parser.parse_args()

    print(args)

    img_path = Path(args.image)
    if args.labels:
        lbl_path = Path(args.labels)
    else:
        lbl_path = img_path.with_suffix(".txt")
        lbl_path = Path(str(lbl_path).replace("images", "labels"))

    if not img_path.exists():
        print("Image not found:", img_path)
        sys.exit(1)
    if not lbl_path.exists():
        print("Label file not found:", lbl_path)
        sys.exit(1)

    # img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
    img = (Image(img_path).get_qt_rgb() * 255).astype(np.uint8)

    if img is None:
        print("Could not load image:", img_path)
        sys.exit(1)

    labels = load_labels_file(str(lbl_path))
    if not labels:
        print("No labels parsed from", lbl_path)
        # continue and just show image
    out = draw_annotations(img.copy(), labels, alpha=args.alpha, draw_bbox_for_poly=(not args.no_bbox))

    out_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
    if 0:
        plt.imshow(out_rgb.transpose(1, 0, 2))
    else:
        plt.imshow(out_rgb)

    for label in labels:
        lclass, coords = label
        # print(lclass, coords)
        bbox = coords[:4]
        # print("bbox", bbox)
        bbox = np.array(bbox) * np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
        yc, xc, h, w = bbox
        # print("bbox", bbox)

        # polyline = np.array(coords[4:]).reshape(-1, 2) * np.array([img.shape[1], img.shape[0]])
        polyline = np.array(coords).reshape(-1, 2) * np.array([img.shape[1], img.shape[0]])
        # print("pl", coords[4:])
        # print("pl", polyline)

        # Convert BGR -> RGB for matplotlib display
        # out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
        # out_rgb = Image()
        plt.plot(polyline[:, 0], polyline[:, 1], "y", linewidth=2)
        if 0:
            plt.plot(
                [yc - h / 2, yc - h / 2, yc + h / 2, yc + h / 2, yc - h / 2],
                [xc - w / 2, xc + w / 2, xc + w / 2, xc - w / 2, xc - w / 2],
                "r",
                linewidth=2,
            )

    # plt.axis("off")
    plt.title(f"{img_path.name} ({lbl_path.name})")
    plt.show()


if __name__ == "__main__":
    main()