test_single_file.py

#!/usr/bin/env python3
"""
Test script to process just one TIFF file to verify functionality
before processing all files.
"""

import os
import sys
import gc
import psutil
from pathlib import Path

# Add src directory to path
src_dir = Path(__file__).parent / "src"
sys.path.insert(0, str(src_dir))

from multispectral_indices import MultispectralIndices, load_multispectral_bands, save_index_as_tiff
from utils import (find_tiff_files, find_shapefile, validate_tiff_bands, create_output_filename, 
                   ensure_directory_exists, get_tiff_info, check_memory_usage)


def check_memory_usage():
    """Check current memory usage and return percentage."""
    memory = psutil.virtual_memory()
    return memory.percent


def test_single_file():
    """Test processing a single TIFF file."""
    
    # Define paths
    project_root = Path(__file__).parent
    data_dir = project_root / "data"
    outputs_dir = project_root / "outputs"
    indices_dir = outputs_dir / "indices"
    
    print("Testing single file processing...")
    print(f"Project root: {project_root}")
    
    # Ensure output directories exist
    ensure_directory_exists(str(outputs_dir))
    ensure_directory_exists(str(indices_dir))
    
    # Find shapefile
    try:
        shapefile_path = find_shapefile(str(data_dir))
    except FileNotFoundError:
        print("No shapefile found!")
        return False
    
    print(f"Found shapefile: {os.path.basename(shapefile_path)}")
    
    # Find first TIFF file
    tiff_files = find_tiff_files(str(data_dir))
    if not tiff_files:
        print("No TIFF files found!")
        return False
    
    tiff_path = tiff_files[0]  # Process only the first file
    print(f"\nProcessing: {os.path.basename(tiff_path)}")
    
    # Check memory before processing
    memory_percent = check_memory_usage()
    print(f"Initial memory usage: {memory_percent:.1f}%")
    
    # Validate file
    if not validate_tiff_bands(tiff_path):
        print("TIFF validation failed!")
        return False
    
    # Get file info
    info = get_tiff_info(tiff_path)
    if 'error' in info:
        print(f"Error reading file: {info['error']}")
        return False
    
    print(f"File info: {info['bands']} bands, {info['width']}x{info['height']}")
    
    try:
        # Load bands with spatial clipping
        print("Loading multispectral bands (clipped to shapefile extent)...")
        bands, profile = load_multispectral_bands(tiff_path, shapefile_path)
        
        memory_percent = check_memory_usage()
        print(f"Memory after loading bands: {memory_percent:.1f}%")
        
        # Initialize indices calculator
        print("Initializing indices calculator...")
        indices_calc = MultispectralIndices(bands)
        
        # Calculate just one index first (NDVI)
        print("Calculating NDVI...")
        ndvi = indices_calc.ndvi()
        
        memory_percent = check_memory_usage()
        print(f"Memory after calculating NDVI: {memory_percent:.1f}%")
        
        # Save NDVI
        output_path = create_output_filename(tiff_path, "NDVI", str(indices_dir))
        save_index_as_tiff(ndvi, output_path, profile)
        print(f"✓ Saved NDVI -> {os.path.basename(output_path)}")
        
        # Clean up
        del bands, indices_calc, ndvi
        gc.collect()
        
        memory_percent = check_memory_usage()
        print(f"Memory after cleanup: {memory_percent:.1f}%")
        
        print("✓ Single file test completed successfully!")
        return True
        
    except Exception as e:
        print(f"Error processing file: {e}")
        gc.collect()
        return False


if __name__ == "__main__":
    success = test_single_file()
    if success:
        print("\n✓ Test passed! The system can handle the TIFF files.")
        print("You can now run the full processing script.")
    else:
        print("\n✗ Test failed! Check the error messages above.")
    
    sys.exit(0 if success else 1)