lang_sam_batch_inference.py

#!/usr/bin/env python3

import os
import argparse
import glob
import tqdm
from PIL import Image
import numpy as np
import cv2
import json
from pathlib import Path

# Import LangSAM
from lang_sam import LangSAM


def _as_numpy_array(arr):
    """Convert torch/Tensor-like arrays to numpy if needed; pass through numpy and lists."""
    if arr is None:
        return None
    if hasattr(arr, 'cpu'):
        arr = arr.cpu().numpy()
        return arr
    if hasattr(arr, 'numpy'):
        arr = arr.numpy()
        return arr
    return arr


def _extract_score_array(result: dict, key: str):
    """Fetch a score array from result and convert to numpy if needed."""
    return _as_numpy_array(result.get(key, None))


def _safe_score_at(arr, idx: int):
    """Return float(arr[idx]) if available; otherwise None."""
    arr = _as_numpy_array(arr)
    if arr is None:
        return None
    try:
        arr = np.atleast_1d(arr)
        if arr.size > idx:
            return float(arr.flat[idx])
    except Exception:
        return None
    return None


def create_visualization(
    image_np,
    mask_binary,
    image_name,
    i,
    visualizations_dir,
    text_prompt,
    alpha=0.4,
    mask_score=None,
    gdino_score=None,
    iou=None,
    wtd_score=None,
):
    """Create 3-panel visualization using OpenCV operations"""
    h, w = image_np.shape[:2]
    
    # Panel 1: Original image
    panel1 = image_np.copy()
    
    # Panel 2: Colored mask
    panel2 = np.zeros_like(image_np)
    panel2[mask_binary] = [0, 255, 0]  # Green for mask regions
    
    # Panel 3: Overlay
    panel3 = image_np.copy()
    colored_mask = np.zeros_like(image_np)
    colored_mask[mask_binary] = [0, 255, 0]
    panel3 = cv2.addWeighted(panel3, 1-alpha, colored_mask, alpha, 0)
    
    # Add text labels to each panel
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = max(0.5, min(w / 800, h / 600))  # Scale font based on image size
    font_thickness = max(1, int(font_scale * 2))
    text_color = (255, 255, 255)  # White text
    bg_color = (0, 0, 0)  # Black background for text
    
    # Add titles with background rectangles and score information
    titles = ['Original Image', 'Segmentation Mask', 'Overlay']
    panels = [panel1, panel2, panel3]
    
    for idx, (panel, title) in enumerate(zip(panels, titles)):
        # Calculate text size for title
        (text_w, text_h), baseline = cv2.getTextSize(title, font, font_scale, font_thickness)
        
        # Prepare score text for mask and overlay panels
        score_text = ""
        if idx == 1 and mask_score is not None:  # Mask panel
            score_text = f"Mask Score: {mask_score:.3f}"
        elif idx == 2:  # Overlay panel
            parts = []
            if gdino_score is not None:
                parts.append(f"GDINO: {gdino_score:.3f}")
            if iou is not None:
                parts.append(f"IoU: {iou:.3f}")
            if wtd_score is not None:
                parts.append(f"Weighted: {wtd_score:.3f}")
            score_text = " | ".join(parts)
        
        # Calculate dimensions for both title and score
        max_width = text_w
        total_height = text_h
        
        if score_text:
            (score_w, score_h), _ = cv2.getTextSize(score_text, font, font_scale * 0.8, font_thickness)
            max_width = max(max_width, score_w)
            total_height += score_h + 5  # Add spacing
        
        # Add background rectangle
        cv2.rectangle(panel, (5, 5), (max_width + 15, total_height + 20), bg_color, -1)
        
        # Add title
        cv2.putText(panel, title, (10, text_h + 10), font, font_scale, text_color, font_thickness)
        
        # Add score text if available
        if score_text:
            cv2.putText(panel, score_text, (10, text_h + score_h + 15), font, font_scale * 0.8, 
                       (255, 255, 0), max(1, font_thickness - 1))  # Yellow color for scores
    
    # Create separator lines (white vertical lines)
    separator_width = 3
    separator = np.ones((h, separator_width, 3), dtype=np.uint8) * 255
    
    # Concatenate panels horizontally with separators
    combined = np.concatenate([panel1, separator, panel2, separator, panel3], axis=1)
    
    # Add overall title at the top
    title_height = 60
    title_bg = np.zeros((title_height, combined.shape[1], 3), dtype=np.uint8)
    
    # Add main title with scores
    main_title = f'LangSAM Results: "{text_prompt}" - {image_name}'
    score_info = []
    if gdino_score is not None:
        score_info.append(f"GDINO: {gdino_score:.3f}")
    if iou is not None:
        score_info.append(f"IoU: {iou:.3f}")
    if mask_score is not None:
        score_info.append(f"Mask: {mask_score:.3f}")
    if wtd_score is not None:
        score_info.append(f"Weighted: {wtd_score:.3f}")
    if score_info:
        main_title += f' [{", ".join(score_info)}]'
    
    title_font_scale = max(0.7, min(combined.shape[1] / 1200, 1.2))
    title_thickness = max(1, int(title_font_scale * 3))
    
    (title_w, title_h), _ = cv2.getTextSize(main_title, font, title_font_scale, title_thickness)
    title_x = max(10, (combined.shape[1] - title_w) // 2)
    title_y = (title_height + title_h) // 2
    
    cv2.putText(title_bg, main_title, (title_x, title_y), font, title_font_scale, 
                (255, 255, 255), title_thickness)
    
    # Combine title and panels
    final_image = np.concatenate([title_bg, combined], axis=0)
    
    # Save the combined visualization
    vis_path = os.path.join(visualizations_dir, f'{image_name}_vis.png')
    cv2.imwrite(vis_path, cv2.cvtColor(final_image, cv2.COLOR_RGB2BGR))


def process_images_with_langsam(model, args):
    """Process all images in input directory with LangSAM model"""
    
    input_dir = args.input_dir
    output_dir = args.output_dir
    text_prompt = args.text_prompt
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Create subdirectories for different outputs
    masks_dir = os.path.join(output_dir, 'masks')
    visualizations_dir = os.path.join(output_dir, 'visualizations')
    annotations_dir = os.path.join(output_dir, 'annotations')
    
    os.makedirs(masks_dir, exist_ok=True)
    os.makedirs(visualizations_dir, exist_ok=True)
    os.makedirs(annotations_dir, exist_ok=True)
    
    # Find all image files in input directory
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff', '*.tif']
    image_files = []
    
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(input_dir, ext)))
        image_files.extend(glob.glob(os.path.join(input_dir, ext.upper())))
    
    if not image_files:
        print(f"No image files found in {input_dir}")
        return
    
    print(f"Found {len(image_files)} images to process")
    print(f"Using text prompt: '{text_prompt}'")
    
    # Process each image
    for image_path in tqdm.tqdm(image_files, desc="Processing images"):
        # Get image filename without extension for output naming
        image_name = Path(image_path).stem
        
        # Load image
        image_pil = Image.open(image_path).convert("RGB")
        image_np = np.array(image_pil)
        
        # Run LangSAM prediction (optionally using external masks for point prompts)
        results = model.predict(
            [image_pil],
            [text_prompt],
            image_path=image_path,
            mask_dir=args.mask_dir if hasattr(args, 'mask_dir') else None
        )
        
        # Process results for this image
        if results and len(results) > 0:
            result = results[0]  # Get result for the first (and only) image
            
            # Extract masks
            masks = result['masks']

            # Extract score components via helpers
            gdino_scores = _extract_score_array(result, 'scores')
            mask_scores = _extract_score_array(result, 'mask_scores')
            ious = _extract_score_array(result, 'ious')
            wtd_scores = _extract_score_array(result, 'wtd_scores')

            # Find the mask with highest weighted score (fallback to GDINO if absent)
            if wtd_scores is not None and np.size(wtd_scores) > 0:
                best_idx = int(np.argmax(wtd_scores))
                # best_idx = int(np.argmax(gdino_scores))
                best_mask = masks[best_idx]
                best_weighted = float(wtd_scores[best_idx])
            elif gdino_scores is not None and np.size(gdino_scores) > 0:
                best_idx = int(np.argmax(gdino_scores))
                best_mask = masks[best_idx]
                best_weighted = None
            else:
                # If no scores available, take the first mask
                best_idx = 0
                best_mask = masks[0] if len(masks) > 0 else None
                best_weighted = None
            
            if best_mask is not None:
                # Convert mask to numpy format
                if hasattr(best_mask, 'cpu'):
                    mask_np = best_mask.cpu().numpy()
                elif hasattr(best_mask, 'numpy'):
                    mask_np = best_mask.numpy()
                else:
                    mask_np = np.array(best_mask)
                    
                # Ensure mask is 2D and binary
                if mask_np.ndim > 2:
                    mask_np = mask_np.squeeze()
                mask_binary = (mask_np > 0.5).astype(bool)
                mask_uint8 = (mask_binary * 255).astype(np.uint8)
                
                # Save the best mask as PNG
                mask_path = os.path.join(masks_dir, f'{image_name}_mask.png')
                cv2.imwrite(mask_path, mask_uint8)
                
                # Prepare score breakdown for visualization
                best_gdino = _safe_score_at(gdino_scores, best_idx)
                best_iou = _safe_score_at(ious, best_idx)
                best_mask_score = _safe_score_at(mask_scores, best_idx)

                # Create 3-panel visualization with detailed scores
                create_visualization(
                    image_np,
                    mask_binary,
                    image_name,
                    0,
                    visualizations_dir,
                    text_prompt,
                    args.vis_alpha,
                    mask_score=best_mask_score,
                    gdino_score=best_gdino,
                    iou=best_iou,
                    wtd_score=best_weighted,
                )
            
            # Save annotation data as JSON (only for the best mask)
            has_mask = 'masks' in result and len(result['masks']) > 0
            annotation_data = {
                'image_path': image_path,
                'image_name': image_name,
                'text_prompt': text_prompt,
                'num_detections': 1 if has_mask else 0,
                'best_mask_index': int(best_idx) if has_mask and 'best_idx' in locals() else None
            }
            
            # Add best mask's bounding box if available
            if 'boxes' in result and has_mask and 'best_idx' in locals():
                boxes = _as_numpy_array(result['boxes'])
                # Convert to regular Python list with float values
                annotation_data['best_box'] = [float(x) for x in boxes[best_idx]]
            
            # Add best mask's logit if available
            if 'logits' in result and has_mask and 'best_idx' in locals():
                logits = _as_numpy_array(result['logits'])
                # Convert to regular Python list with float values
                annotation_data['best_logit'] = [float(x) for x in logits[best_idx]]
            
            # Add weighted and component scores if available
            if best_weighted is not None:
                annotation_data['best_weighted_score'] = float(best_weighted)
            best_gdino_for_json = _safe_score_at(gdino_scores, best_idx)
            if best_gdino_for_json is not None:
                annotation_data['best_gdino_score'] = best_gdino_for_json
            best_iou_for_json = _safe_score_at(ious, best_idx)
            if best_iou_for_json is not None:
                annotation_data['best_iou'] = best_iou_for_json
            best_mask_for_json = _safe_score_at(mask_scores, best_idx)
            if best_mask_for_json is not None:
                annotation_data['best_mask_score'] = best_mask_for_json
            
            # Save annotation JSON
            annotation_path = os.path.join(annotations_dir, f'{image_name}_annotation.json')
            with open(annotation_path, 'w') as f:
                json.dump(annotation_data, f, indent=2)
            
            score_info = ""
            if best_weighted is not None:
                score_info = f" (weighted: {best_weighted:.3f})"
            
            # print(f"Processed {image_name}: found {annotation_data['num_detections']} detection{score_info}")
        
        else:
            print(f"No results found for {image_name}")
                
    print(f"\nProcessing complete! Results saved to {output_dir}")
    print(f"- Best masks saved to: {masks_dir}")
    print(f"  * Files ending with '_mask.png' contain the highest-scoring mask for each image")
    print(f"- 3-panel visualizations saved to: {visualizations_dir}")
    print(f"  * Files ending with '_vis.png' show original, best mask, and overlay side by side")
    print(f"- Annotations saved to: {annotations_dir}")
    print(f"  * JSON files contain metadata for the best detection per image")


def main():
    parser = argparse.ArgumentParser(description='Run LangSAM inference on all images in a directory')
    parser.add_argument('--input_dir', required=True, type=str,
                       help='Input directory containing images to process')
    parser.add_argument('--output_dir', required=True, type=str,
                       help='Output directory to save results')
    parser.add_argument('--text_prompt', default='sky', type=str,
                       help='Text prompt for segmentation (default: "sky")')
    parser.add_argument('--vis_alpha', default=0.4, type=float,
                       help='Alpha blending value for overlay visualization (default: 0.4)')
    parser.add_argument('--mask_dir', default=None, type=str,
                       help='Directory containing binary masks to use as point prompts')
    
    args = parser.parse_args()
    
    # Validate input directory
    if not os.path.exists(args.input_dir):
        raise ValueError(f"Input directory does not exist: {args.input_dir}")
    
    if not os.path.isdir(args.input_dir):
        raise ValueError(f"Input path is not a directory: {args.input_dir}")
    
    print(f"Input directory: {args.input_dir}")
    print(f"Output directory: {args.output_dir}")
    print(f"Text prompt: {args.text_prompt}")
    print(f"Overlay alpha: {args.vis_alpha}")
    
    # Initialize LangSAM model
    print("Loading LangSAM model...")
    model = LangSAM()
    print("Model loaded successfully!")
    
    # Process images
    process_images_with_langsam(model, args)


if __name__ == '__main__':
    main()