calculate_organ_areas_by_class.py

"""
Calculate organ-specific areas from UltraSam model predictions
Extracts per-class segmentation areas for each organ type
"""

import os
os.environ['MPLBACKEND'] = 'Agg'
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '0'
os.environ['QT_QPA_PLATFORM'] = 'offscreen'

import sys
sys.path.insert(0, '.')

import torch
import numpy as np
from PIL import Image
from pathlib import Path
import argparse
import pandas as pd
import json
import random
from tqdm import tqdm

from mmengine.config import Config
from mmdet.apis import init_detector, inference_detector


# Organ class names from FASS dataset
ORGAN_CLASSES = ['artery', 'liver', 'stomach', 'vein']


def calculate_organ_areas_from_predictions(model, img_path, pixel_to_mm2=None):
    """
    Run model inference and calculate area for each organ class
    
    Args:
        model: Trained UltraSam model
        img_path: Path to ultrasound image
        pixel_to_mm2: Conversion factor from pixels to mm² (optional)
    
    Returns:
        Dictionary with per-organ areas
    """
    # Run inference
    result = inference_detector(model, str(img_path))
    
    # Extract predictions
    pred_instances = result.pred_instances
    
    if len(pred_instances) == 0:
        # No predictions
        return {organ: 0 for organ in ORGAN_CLASSES}
    
    # Get masks and labels
    masks = pred_instances.masks.cpu().numpy()  # [N, H, W]
    labels = pred_instances.labels.cpu().numpy()  # [N]
    scores = pred_instances.scores.cpu().numpy()  # [N]
    
    # Initialize areas for each class
    organ_areas = {organ: 0 for organ in ORGAN_CLASSES}
    organ_counts = {organ: 0 for organ in ORGAN_CLASSES}
    
    # Calculate area for each prediction
    for mask, label, score in zip(masks, labels, scores):
        organ_name = ORGAN_CLASSES[label]
        area_pixels = np.sum(mask)
        
        organ_areas[organ_name] += area_pixels
        organ_counts[organ_name] += 1
    
    # Convert to mm² if conversion factor provided
    if pixel_to_mm2 is not None:
        organ_areas_mm2 = {k: v * pixel_to_mm2 for k, v in organ_areas.items()}
    else:
        organ_areas_mm2 = None
    
    return {
        'areas_pixels': organ_areas,
        'areas_mm2': organ_areas_mm2,
        'counts': organ_counts,
        'total_predictions': len(pred_instances)
    }


def analyze_dataset_organ_areas(
    config_path,
    checkpoint_path,
    image_dir=None,
    coco_json=None,
    output_csv='organ_areas_by_class.csv',
    n_images=None,
    seed=42,
    pixel_to_mm2=None
):
    """
    Analyze organ areas across entire dataset
    
    Args:
        config_path: Path to model config
        checkpoint_path: Path to model checkpoint
        image_dir: Directory with images (if not using COCO)
        coco_json: Path to COCO annotation file (optional)
        output_csv: Output CSV path
        n_images: Number of images to analyze (None = all)
        seed: Random seed
        pixel_to_mm2: Conversion factor to mm²
    """
    print("="*70)
    print("Organ-Specific Area Calculation (Per-Class Segmentation)")
    print("="*70)
    
    # Load model
    print(f"\n📦 Loading UltraSam model...")
    print(f"   Config: {config_path}")
    print(f"   Checkpoint: {checkpoint_path}")
    
    cfg = Config.fromfile(config_path)
    model = init_detector(cfg, checkpoint_path, device='cuda:0')
    model.eval()
    
    print("✓ Model loaded successfully")
    print(f"   Organ classes: {', '.join(ORGAN_CLASSES)}")
    
    # Get image list
    if coco_json:
        print(f"\n📂 Loading images from COCO: {coco_json}")
        with open(coco_json, 'r') as f:
            coco_data = json.load(f)
        
        image_root = Path(cfg.test_dataloader.dataset.get('data_root', ''))
        img_prefix = cfg.test_dataloader.dataset.get('data_prefix', {}).get('img', '')
        
        all_images = []
        for img_info in coco_data['images']:
            img_path = image_root / img_prefix / img_info['file_name']
            all_images.append((str(img_path), img_info['id'], img_info['file_name']))
    
    elif image_dir:
        print(f"\n📂 Loading images from directory: {image_dir}")
        image_dir = Path(image_dir)
        all_images = [(str(p), p.stem, p.name) 
                      for p in sorted(image_dir.glob('*.png')) + sorted(image_dir.glob('*.jpg'))]
    else:
        raise ValueError("Must provide either image_dir or coco_json")
    
    # Sample if requested
    if n_images and n_images < len(all_images):
        random.seed(seed)
        selected_images = random.sample(all_images, n_images)
        selected_images = sorted(selected_images, key=lambda x: x[2])
    else:
        selected_images = all_images
    
    print(f"📊 Processing {len(selected_images)} images")
    if pixel_to_mm2:
        print(f"🔬 Pixel to mm² conversion: {pixel_to_mm2}")
    print(f"\n{'='*70}\n")
    
    # Process images
    results = []
    
    for img_path, img_id, img_name in tqdm(selected_images, desc="Analyzing organs"):
        try:
            # Calculate organ areas
            organ_data = calculate_organ_areas_from_predictions(model, img_path, pixel_to_mm2)
            
            # Build result row
            result = {
                'image_name': img_name,
                'image_id': img_id,
                'total_predictions': organ_data['total_predictions']
            }
            
            # Add per-organ data
            for organ in ORGAN_CLASSES:
                result[f'{organ}_area_pixels'] = organ_data['areas_pixels'][organ]
                result[f'{organ}_count'] = organ_data['counts'][organ]
                
                if organ_data['areas_mm2']:
                    result[f'{organ}_area_mm2'] = organ_data['areas_mm2'][organ]
            
            # Total area
            result['total_area_pixels'] = sum(organ_data['areas_pixels'].values())
            if organ_data['areas_mm2']:
                result['total_area_mm2'] = sum(organ_data['areas_mm2'].values())
            
            results.append(result)
            
        except Exception as e:
            print(f"\n✗ Error processing {img_name}: {e}")
            continue
    
    # Create DataFrame
    df = pd.DataFrame(results)
    
    # Calculate statistics
    print(f"\n{'='*70}")
    print(f"📈 Summary Statistics")
    print(f"{'='*70}")
    print(f"\nTotal images analyzed: {len(results)}")
    print(f"Total predictions: {df['total_predictions'].sum():,}")
    
    print(f"\n{'Per-Organ Statistics (Pixels)':}")
    print(f"{'─'*70}")
    
    for organ in ORGAN_CLASSES:
        organ_col = f'{organ}_area_pixels'
        count_col = f'{organ}_count'
        
        if organ_col in df.columns:
            total_area = df[organ_col].sum()
            mean_area = df[organ_col].mean()
            median_area = df[organ_col].median()
            std_area = df[organ_col].std()
            total_count = df[count_col].sum()
            images_with_organ = (df[organ_col] > 0).sum()
            
            print(f"\n{organ.upper()}:")
            print(f"  Images with {organ}: {images_with_organ}/{len(df)} ({images_with_organ/len(df)*100:.1f}%)")
            print(f"  Total instances: {total_count}")
            print(f"  Total area: {total_area:,.0f} pixels")
            print(f"  Mean area per image: {mean_area:,.0f} ± {std_area:,.0f} pixels")
            print(f"  Median area: {median_area:,.0f} pixels")
    
    # Save to CSV
    output_path = Path(output_csv)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_path, index=False)
    
    print(f"\n{'='*70}")
    print(f"💾 Results saved to: {output_path.absolute()}")
    print(f"{'='*70}")
    
    # Show top images by organ
    print(f"\n{'Top 5 Images by Organ Area':}")
    print(f"{'─'*70}")
    
    for organ in ORGAN_CLASSES:
        organ_col = f'{organ}_area_pixels'
        if organ_col in df.columns and df[organ_col].sum() > 0:
            top_5 = df.nlargest(5, organ_col)[['image_name', organ_col, f'{organ}_count']]
            print(f"\n{organ.upper()}:")
            for idx, row in top_5.iterrows():
                if row[organ_col] > 0:
                    print(f"  {row['image_name']}: {row[organ_col]:,.0f} px ({row[f'{organ}_count']} instances)")
    
    print(f"\n{'='*70}")
    
    return df


def main():
    parser = argparse.ArgumentParser(
        description='Calculate organ-specific areas using UltraSam model predictions'
    )
    parser.add_argument(
        '--config',
        type=str,
        default='work_dir/boundary_loss_exp/UltraSAM_box_refine.py',
        help='Path to model config'
    )
    parser.add_argument(
        '--checkpoint',
        type=str,
        default='work_dir/boundary_loss_exp/iter_30000.pth',
        help='Path to model checkpoint'
    )
    parser.add_argument(
        '--image-dir',
        type=str,
        default='work_dir/datasets/images',
        help='Directory containing images'
    )
    parser.add_argument(
        '--coco-json',
        type=str,
        default=None,
        help='Path to COCO annotation file (optional)'
    )
    parser.add_argument(
        '--output-csv',
        type=str,
        default='work_dir/boundary_loss_inference/20251125_053719/work_dir/organ_areas_by_class.csv',
        help='Output CSV path'
    )
    parser.add_argument(
        '--n-images',
        type=int,
        default=None,
        help='Number of images to analyze (None = all)'
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed'
    )
    parser.add_argument(
        '--pixel-to-mm2',
        type=float,
        default=None,
        help='Conversion factor from pixels to mm² (optional)'
    )
    
    args = parser.parse_args()
    
    df = analyze_dataset_organ_areas(
        args.config,
        args.checkpoint,
        image_dir=args.image_dir,
        coco_json=args.coco_json,
        output_csv=args.output_csv,
        n_images=args.n_images,
        seed=args.seed,
        pixel_to_mm2=args.pixel_to_mm2
    )
    
    print(f"\n✅ Analysis complete!")
    print(f"📄 Detailed per-organ results saved to: {args.output_csv}")


if __name__ == '__main__':
    main()