64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for memory-efficient batch loader.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add src directory to path for imports
|
|
sys.path.append(str(Path(__file__).parent / "src"))
|
|
|
|
from src.batch_loader import MemoryEfficientBatchLoader, BatchProcessor
|
|
from src.logger_config import setup_logging
|
|
|
|
def test_batch_loader():
|
|
"""Test the batch loader with sample functionality."""
|
|
|
|
# Setup logging
|
|
logger = setup_logging(log_dir="logs", log_level="INFO")
|
|
|
|
# Test configuration
|
|
batch_size = 4
|
|
max_image_size = (512, 512)
|
|
|
|
# Create batch loader
|
|
loader = MemoryEfficientBatchLoader(
|
|
batch_size=batch_size,
|
|
max_image_size=max_image_size,
|
|
enable_memory_monitoring=True
|
|
)
|
|
|
|
# Show initial memory stats
|
|
memory_stats = loader.get_memory_stats()
|
|
logger.info(f"Initial memory stats: {memory_stats}")
|
|
|
|
# Sample image paths (for testing - these would be real image paths)
|
|
sample_image_paths = [
|
|
f"/home/nguyendc/thanh-dev/IQA-Metric-Benchmark/data/task/cni/images/image_{i}.png"
|
|
for i in range(10)
|
|
]
|
|
|
|
# Test batch creation
|
|
batches = list(loader.create_batches(sample_image_paths))
|
|
logger.info(f"Created {len(batches)} batches from {len(sample_image_paths)} images")
|
|
for i, batch in enumerate(batches):
|
|
logger.info(f"Batch {i+1}: {len(batch)} images")
|
|
|
|
# Mock scoring function for testing
|
|
def mock_scoring_function(images):
|
|
"""Mock function that returns random scores."""
|
|
import random
|
|
return [random.uniform(1.0, 5.0) for _ in images]
|
|
|
|
# Test processing a single batch (with mock data)
|
|
if batches:
|
|
logger.info("Testing batch processing with mock data...")
|
|
# Note: This will fail with actual file paths that don't exist
|
|
# but demonstrates the interface
|
|
|
|
logger.info("Batch loader test completed!")
|
|
|
|
if __name__ == "__main__":
|
|
test_batch_loader()
|