scripts/validate_dataset.py

#!/usr/bin/env python3
"""
Dataset Validation Script for Flux Identity LoRA Training

Features:
- Scan all images in dataset directory
- Detect corrupt/unreadable files using PIL
- Report image resolutions
- Find duplicate files using perceptual hashing (imagehash)
- Verify matching caption files (.txt) for each image
- Generate validation report
"""

import os
import sys
import argparse
from pathlib import Path
from collections import defaultdict
from PIL import Image
import imagehash
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn

console = Console()

# Supported image formats
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff', '.tif'}


def scan_images(directory: Path) -> list:
    """Scan directory for image files."""
    images = []
    for ext in IMAGE_EXTENSIONS:
        images.extend(directory.glob(f'*{ext}'))
        images.extend(directory.glob(f'*{ext.upper()}'))
    return sorted(set(images))


def validate_image(image_path: Path) -> dict:
    """Validate a single image file."""
    result = {
        'path': image_path,
        'valid': False,
        'width': 0,
        'height': 0,
        'format': None,
        'mode': None,
        'error': None,
        'phash': None,
        'has_caption': False,
        'caption_path': None
    }

    try:
        with Image.open(image_path) as img:
            img.verify()

        # Reopen for actual data (verify() can leave file in bad state)
        with Image.open(image_path) as img:
            result['valid'] = True
            result['width'] = img.width
            result['height'] = img.height
            result['format'] = img.format
            result['mode'] = img.mode

            # Calculate perceptual hash for duplicate detection
            result['phash'] = str(imagehash.phash(img))

    except Exception as e:
        result['error'] = str(e)

    # Check for caption file
    caption_path = image_path.with_suffix('.txt')
    if caption_path.exists():
        result['has_caption'] = True
        result['caption_path'] = caption_path

    return result


def find_duplicates(results: list) -> dict:
    """Find duplicate images using perceptual hashing."""
    hash_groups = defaultdict(list)

    for r in results:
        if r['phash']:
            hash_groups[r['phash']].append(r['path'])

    # Return only groups with duplicates
    return {h: paths for h, paths in hash_groups.items() if len(paths) > 1}


def generate_report(results: list, duplicates: dict, output_file: Path = None):
    """Generate validation report."""

    valid_count = sum(1 for r in results if r['valid'])
    invalid_count = sum(1 for r in results if not r['valid'])
    with_caption = sum(1 for r in results if r['has_caption'])
    without_caption = sum(1 for r in results if r['valid'] and not r['has_caption'])

    console.print("\n[bold blue]═══ Dataset Validation Report ═══[/bold blue]\n")

    # Summary table
    summary = Table(title="Summary", show_header=False)
    summary.add_column("Metric", style="cyan")
    summary.add_column("Value", style="green")

    summary.add_row("Total Images Scanned", str(len(results)))
    summary.add_row("Valid Images", str(valid_count))
    summary.add_row("Invalid/Corrupt Images", str(invalid_count))
    summary.add_row("Images with Captions", str(with_caption))
    summary.add_row("Images Missing Captions", str(without_caption))
    summary.add_row("Duplicate Groups Found", str(len(duplicates)))

    console.print(summary)

    # Resolution distribution
    resolutions = defaultdict(int)
    for r in results:
        if r['valid']:
            res = f"{r['width']}x{r['height']}"
            resolutions[res] += 1

    if resolutions:
        console.print("\n[bold]Resolution Distribution:[/bold]")
        res_table = Table()
        res_table.add_column("Resolution", style="cyan")
        res_table.add_column("Count", style="green")
        res_table.add_column("Percentage", style="yellow")

        for res, count in sorted(resolutions.items(), key=lambda x: -x[1])[:10]:
            pct = (count / valid_count) * 100 if valid_count > 0 else 0
            res_table.add_row(res, str(count), f"{pct:.1f}%")

        console.print(res_table)

    # Invalid files
    invalid_files = [r for r in results if not r['valid']]
    if invalid_files:
        console.print("\n[bold red]Invalid/Corrupt Files:[/bold red]")
        for r in invalid_files:
            console.print(f"  • {r['path'].name}: {r['error']}")

    # Missing captions
    missing_captions = [r for r in results if r['valid'] and not r['has_caption']]
    if missing_captions:
        console.print("\n[bold yellow]Images Missing Captions:[/bold yellow]")
        for r in missing_captions[:20]:  # Show first 20
            console.print(f"  • {r['path'].name}")
        if len(missing_captions) > 20:
            console.print(f"  ... and {len(missing_captions) - 20} more")

    # Duplicates
    if duplicates:
        console.print("\n[bold yellow]Potential Duplicates (by perceptual hash):[/bold yellow]")
        for hash_val, paths in list(duplicates.items())[:10]:
            console.print(f"\n  Hash: {hash_val}")
            for p in paths:
                console.print(f"    • {p.name}")
        if len(duplicates) > 10:
            console.print(f"\n  ... and {len(duplicates) - 10} more duplicate groups")

    # Recommendations
    console.print("\n[bold blue]Recommendations:[/bold blue]")

    if invalid_count > 0:
        console.print(f"  ⚠ Remove or fix {invalid_count} corrupt image(s)")

    if without_caption > 0:
        console.print(f"  ⚠ Add captions for {without_caption} image(s)")

    if duplicates:
        total_dups = sum(len(paths) - 1 for paths in duplicates.values())
        console.print(f"  ⚠ Review {total_dups} potential duplicate image(s)")

    # Check for resolution issues
    small_images = [r for r in results if r['valid'] and (r['width'] < 512 or r['height'] < 512)]
    if small_images:
        console.print(f"  ⚠ {len(small_images)} image(s) are smaller than 512px (may affect quality)")

    if valid_count > 0 and without_caption == 0 and invalid_count == 0 and not duplicates:
        console.print("  ✓ Dataset looks good! Ready for training.")

    # Save report to file if requested
    if output_file:
        with open(output_file, 'w') as f:
            f.write("Dataset Validation Report\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"Total Images: {len(results)}\n")
            f.write(f"Valid: {valid_count}\n")
            f.write(f"Invalid: {invalid_count}\n")
            f.write(f"With Captions: {with_caption}\n")
            f.write(f"Missing Captions: {without_caption}\n")
            f.write(f"Duplicate Groups: {len(duplicates)}\n")

            if invalid_files:
                f.write("\nInvalid Files:\n")
                for r in invalid_files:
                    f.write(f"  - {r['path']}: {r['error']}\n")

            if missing_captions:
                f.write("\nMissing Captions:\n")
                for r in missing_captions:
                    f.write(f"  - {r['path']}\n")

        console.print(f"\n[dim]Report saved to: {output_file}[/dim]")


def main():
    parser = argparse.ArgumentParser(description="Validate dataset for Flux LoRA training")
    parser.add_argument(
        "directory",
        nargs="?",
        default="/workspace/flux-project/datasets/identity/images",
        help="Directory containing training images"
    )
    parser.add_argument(
        "-o", "--output",
        help="Save report to file"
    )
    parser.add_argument(
        "--no-duplicates",
        action="store_true",
        help="Skip duplicate detection (faster)"
    )

    args = parser.parse_args()

    dataset_dir = Path(args.directory)

    if not dataset_dir.exists():
        console.print(f"[red]Error: Directory not found: {dataset_dir}[/red]")
        sys.exit(1)

    console.print(f"[bold]Scanning directory: {dataset_dir}[/bold]")

    # Find all images
    images = scan_images(dataset_dir)

    if not images:
        console.print("[yellow]No images found in directory.[/yellow]")
        sys.exit(0)

    console.print(f"Found {len(images)} image file(s)\n")

    # Validate each image
    results = []
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console
    ) as progress:
        task = progress.add_task("Validating images...", total=len(images))

        for img_path in images:
            result = validate_image(img_path)
            results.append(result)
            progress.advance(task)

    # Find duplicates
    duplicates = {}
    if not args.no_duplicates:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Checking for duplicates...", total=1)
            duplicates = find_duplicates(results)
            progress.advance(task)

    # Generate report
    output_file = Path(args.output) if args.output else None
    generate_report(results, duplicates, output_file)


if __name__ == "__main__":
    main()