#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Visualize YOLOv8-seg labels by drawing polygon masks on images. This helps verify that label conversion (e.g., from Labelme JSON to YOLO .txt) is correct. Example: python3 segmentation/visualize_yolo_seg_labels.py \ --dataset ./datasets/fish_body_seg_filtered \ --output ./visualizations/yolo_seg_labels \ --max_images 50 \ --split train """ from __future__ import annotations import argparse import random from pathlib import Path from typing import List, Tuple import cv2 import numpy as np def parse_yolo_seg_label(label_path: Path, img_w: int, img_h: int) -> List[np.ndarray]: """ Parse YOLO segmentation label file. Returns list of polygons (each as Nx2 numpy array in pixel coordinates). """ if not label_path.exists(): return [] polygons: List[np.ndarray] = [] lines = label_path.read_text(encoding="utf-8").strip().split("\n") for line in lines: line = line.strip() if not line: continue parts = line.split() if len(parts) < 7: # class_id + at least 3 points (x,y pairs) continue try: _class_id = int(parts[0]) coords = [float(x) for x in parts[1:]] if len(coords) % 2 != 0: continue # Convert normalized [0,1] to pixel coordinates points = [] for i in range(0, len(coords), 2): x_norm = coords[i] y_norm = coords[i + 1] x_px = int(x_norm * img_w) y_px = int(y_norm * img_h) points.append([x_px, y_px]) if len(points) >= 3: polygons.append(np.array(points, dtype=np.int32)) except (ValueError, IndexError): continue return polygons def draw_polygons_on_image( img: np.ndarray, polygons: List[np.ndarray], class_colors: List[Tuple[int, int, int]], alpha: float = 0.5 ) -> np.ndarray: """ Draw polygons as semi-transparent masks on image. Returns a new image with overlays. """ overlay = img.copy() mask = np.zeros(img.shape[:2], dtype=np.uint8) for i, poly in enumerate(polygons): color_idx = i % len(class_colors) color = class_colors[color_idx] cv2.fillPoly(mask, [poly], 255) cv2.fillPoly(overlay, [poly], color) # Also draw outline cv2.polylines(overlay, [poly], isClosed=True, color=color, thickness=2) # Blend overlay with original result = cv2.addWeighted(overlay, alpha, img, 1.0 - alpha, 0) return result def visualize_dataset( dataset_dir: Path, output_dir: Path, split: str = "train", max_images: int = 50, class_names: List[str] = None, alpha: float = 0.5, ) -> None: """ Visualize YOLO segmentation labels on images. """ dataset_dir = dataset_dir.expanduser().resolve() if not dataset_dir.exists(): raise SystemExit(f"dataset_dir not found: {dataset_dir}") img_dir = dataset_dir / "images" / split lbl_dir = dataset_dir / "labels" / split if not img_dir.exists(): raise SystemExit(f"images directory not found: {img_dir}") if not lbl_dir.exists(): raise SystemExit(f"labels directory not found: {lbl_dir}") # Collect image-label pairs pairs: List[Tuple[Path, Path]] = [] for img_path in sorted(img_dir.iterdir()): if img_path.suffix.lower() not in {".jpg", ".jpeg", ".png", ".bmp"}: continue lbl_path = lbl_dir / f"{img_path.stem}.txt" if lbl_path.exists(): pairs.append((img_path, lbl_path)) if not pairs: raise SystemExit(f"No image-label pairs found in {split} split") # Limit number of images if max_images > 0 and len(pairs) > max_images: random.seed(42) pairs = random.sample(pairs, max_images) print(f"Visualizing {len(pairs)} images from {split} split...") # Generate colors for classes (BGR format for OpenCV) if class_names is None: class_names = ["class0", "class1", "class2"] colors = [ (0, 255, 0), # green (255, 0, 0), # blue (0, 0, 255), # red (255, 255, 0), # cyan (255, 0, 255), # magenta (0, 255, 255), # yellow ] class_colors = colors[: len(class_names)] output_dir = output_dir.expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) for img_path, lbl_path in pairs: # Load image img = cv2.imread(str(img_path)) if img is None: print(f"[warn] failed to load: {img_path}") continue h, w = img.shape[:2] # Parse labels polygons = parse_yolo_seg_label(lbl_path, w, h) if not polygons: print(f"[warn] no polygons found in: {lbl_path}") # Still save the original image for reference out_path = output_dir / f"{img_path.stem}_no_labels{img_path.suffix}" cv2.imwrite(str(out_path), img) continue # Draw polygons vis_img = draw_polygons_on_image(img, polygons, class_colors, alpha=alpha) # Save visualization out_path = output_dir / f"{img_path.stem}_vis{img_path.suffix}" cv2.imwrite(str(out_path), vis_img) print(f"[done] saved {len(pairs)} visualizations to: {output_dir}") def main() -> None: parser = argparse.ArgumentParser(description="Visualize YOLOv8-seg labels on images") parser.add_argument("--dataset", type=str, required=True, help="Path to YOLO-seg dataset directory") parser.add_argument("--output", type=str, required=True, help="Output directory for visualizations") parser.add_argument( "--split", type=str, default="train", choices=["train", "val", "test"], help="Dataset split to visualize" ) parser.add_argument("--max_images", type=int, default=50, help="Maximum number of images to visualize (0=all)") parser.add_argument( "--classes", type=str, default="fishbody", help="Comma-separated class names (for color assignment, e.g. 'fishbody' or 'body,fin,tail')", ) parser.add_argument( "--alpha", type=float, default=0.5, help="Transparency of mask overlay (0.0=transparent, 1.0=opaque)" ) args = parser.parse_args() class_names = [c.strip() for c in (args.classes or "").split(",") if c.strip()] if not class_names: class_names = ["class0"] visualize_dataset( dataset_dir=Path(args.dataset), output_dir=Path(args.output), split=args.split, max_images=args.max_images, class_names=class_names, alpha=args.alpha, ) if __name__ == "__main__": main()