282 lines
10 KiB
Python
282 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for IRE Shadow Analysis functionality.
|
|
Demonstrates the false color IRE stops approach for shadow detection in sRGB images.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from PIL import Image
|
|
import sys
|
|
import os
|
|
|
|
# Add the project root to the path
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from utils.ire_shadow_analyzer import IREShadowAnalyzer
|
|
from utils.debug_visualizer import DebugVisualizer
|
|
|
|
def create_test_image():
|
|
"""
|
|
Create a test sRGB image with various shadow characteristics.
|
|
"""
|
|
# Create a 512x512 test image
|
|
height, width = 512, 512
|
|
image = torch.zeros((1, height, width, 3))
|
|
|
|
# Create a gradient from left to right (simulating lighting)
|
|
x_coords = torch.linspace(0, 1, width).unsqueeze(0).repeat(height, 1)
|
|
y_coords = torch.linspace(0, 1, height).unsqueeze(1).repeat(1, width)
|
|
|
|
# Base lighting gradient
|
|
base_light = 0.3 + 0.7 * x_coords # Left side darker, right side brighter
|
|
|
|
# Add some objects with shadows
|
|
# Object 1: Hard shadow (sharp transition)
|
|
obj1_x, obj1_y = 150, 200
|
|
obj1_size = 80
|
|
obj1_mask = ((torch.abs(torch.arange(width) - obj1_x) < obj1_size) &
|
|
(torch.abs(torch.arange(height).unsqueeze(1) - obj1_y) < obj1_size))
|
|
|
|
# Hard shadow: sharp cutoff
|
|
hard_shadow = torch.where(
|
|
(x_coords < 0.4) & (y_coords > 0.3) & (y_coords < 0.7),
|
|
torch.tensor(0.1), # Very dark shadow
|
|
base_light
|
|
)
|
|
|
|
# Object 2: Soft shadow (gradual transition)
|
|
obj2_x, obj2_y = 350, 300
|
|
obj2_size = 60
|
|
|
|
# Soft shadow: gradual transition using Gaussian-like falloff
|
|
shadow_center_x, shadow_center_y = 0.6, 0.5
|
|
shadow_distance = torch.sqrt((x_coords - shadow_center_x)**2 + (y_coords - shadow_center_y)**2)
|
|
soft_shadow_factor = torch.exp(-shadow_distance * 3) # Gradual falloff
|
|
|
|
soft_shadow = base_light * (1 - 0.6 * soft_shadow_factor)
|
|
|
|
# Combine both shadow types
|
|
final_image = torch.minimum(hard_shadow, soft_shadow)
|
|
|
|
# Add some noise for realism
|
|
noise = torch.randn_like(final_image) * 0.02
|
|
final_image = torch.clamp(final_image + noise, 0, 1)
|
|
|
|
# Convert to RGB
|
|
image[0, :, :, 0] = final_image # Red channel
|
|
image[0, :, :, 1] = final_image # Green channel
|
|
image[0, :, :, 2] = final_image # Blue channel
|
|
|
|
return image
|
|
|
|
def test_ire_analysis():
|
|
"""
|
|
Test the IRE shadow analysis functionality.
|
|
"""
|
|
print("=== IRE Shadow Analysis Test ===")
|
|
|
|
# Create test image
|
|
print("Creating test image with hard and soft shadows...")
|
|
test_image = create_test_image()
|
|
|
|
# Initialize IRE analyzer
|
|
analyzer = IREShadowAnalyzer(
|
|
shadow_ire_threshold=20.0,
|
|
transition_sensitivity=0.1
|
|
)
|
|
|
|
# Perform analysis
|
|
print("Performing IRE analysis...")
|
|
results = analyzer.generate_ire_analysis_report(test_image)
|
|
|
|
# Extract results
|
|
ire_values = results['ire_values']
|
|
false_color = results['false_color_visualization']
|
|
transition_analysis = results['transition_analysis']
|
|
shadow_classification = results['shadow_classification']
|
|
ire_stats = results['ire_statistics']
|
|
|
|
# Print results
|
|
print("\n=== IRE Analysis Results ===")
|
|
print(f"Shadow Character: {shadow_classification['shadow_character']}")
|
|
print(f"Transition Quality: {shadow_classification['transition_quality']}")
|
|
print(f"Soft Shadow Ratio: {shadow_classification['soft_ratio']}")
|
|
print(f"Hard Shadow Ratio: {shadow_classification['hard_ratio']}")
|
|
print(f"Mean Gradient: {shadow_classification['mean_gradient']}")
|
|
|
|
print(f"\n=== IRE Statistics ===")
|
|
print(f"Min IRE: {ire_stats['min_ire']:.1f}")
|
|
print(f"Max IRE: {ire_stats['max_ire']:.1f}")
|
|
print(f"Mean IRE: {ire_stats['mean_ire']:.1f}")
|
|
print(f"Shadow Percentage: {ire_stats['shadow_percentage']:.1f}%")
|
|
|
|
# Create visualizations
|
|
print("\nCreating visualizations...")
|
|
|
|
# Convert tensors to numpy for visualization
|
|
original_np = test_image[0].numpy()
|
|
ire_np = ire_values[0].numpy()
|
|
false_color_np = false_color[0].numpy()
|
|
shadow_mask_np = transition_analysis['shadow_mask'][0].numpy()
|
|
soft_shadow_np = transition_analysis['soft_shadow_mask'][0].numpy()
|
|
hard_shadow_np = transition_analysis['hard_shadow_mask'][0].numpy()
|
|
|
|
# Create figure with subplots
|
|
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
|
|
fig.suptitle('IRE Shadow Analysis Results', fontsize=16)
|
|
|
|
# Original image
|
|
axes[0, 0].imshow(original_np)
|
|
axes[0, 0].set_title('Original sRGB Image')
|
|
axes[0, 0].axis('off')
|
|
|
|
# IRE values (grayscale)
|
|
im1 = axes[0, 1].imshow(ire_np, cmap='gray', vmin=0, vmax=100)
|
|
axes[0, 1].set_title('IRE Values (0-100)')
|
|
axes[0, 1].axis('off')
|
|
plt.colorbar(im1, ax=axes[0, 1], fraction=0.046, pad=0.04)
|
|
|
|
# False color visualization
|
|
axes[0, 2].imshow(false_color_np)
|
|
axes[0, 2].set_title('False Color IRE Visualization')
|
|
axes[0, 2].axis('off')
|
|
|
|
# Shadow mask
|
|
axes[1, 0].imshow(shadow_mask_np, cmap='gray')
|
|
axes[1, 0].set_title('Shadow Mask')
|
|
axes[1, 0].axis('off')
|
|
|
|
# Soft shadows
|
|
axes[1, 1].imshow(soft_shadow_np, cmap='Blues')
|
|
axes[1, 1].set_title('Soft Shadow Areas')
|
|
axes[1, 1].axis('off')
|
|
|
|
# Hard shadows
|
|
axes[1, 2].imshow(hard_shadow_np, cmap='Reds')
|
|
axes[1, 2].set_title('Hard Shadow Areas')
|
|
axes[1, 2].axis('off')
|
|
|
|
plt.tight_layout()
|
|
plt.savefig('ire_analysis_results.png', dpi=150, bbox_inches='tight')
|
|
print("Results saved to 'ire_analysis_results.png'")
|
|
|
|
# Create IRE legend
|
|
legend = analyzer.create_ire_legend()
|
|
legend_np = legend[0].numpy()
|
|
|
|
plt.figure(figsize=(8, 4))
|
|
plt.imshow(legend_np)
|
|
plt.title('IRE Color Mapping Legend')
|
|
plt.axis('off')
|
|
|
|
# Add IRE range labels
|
|
ire_ranges = ['0-10', '10-20', '20-30', '30-50', '50-70', '70-90', '90-100']
|
|
y_positions = np.linspace(legend_np.shape[0] - 20, 20, len(ire_ranges))
|
|
for i, (label, y_pos) in enumerate(zip(ire_ranges, y_positions)):
|
|
plt.text(legend_np.shape[1] + 10, y_pos, f'{label} IRE',
|
|
verticalalignment='center', fontsize=10)
|
|
|
|
plt.tight_layout()
|
|
plt.savefig('ire_legend.png', dpi=150, bbox_inches='tight')
|
|
print("Legend saved to 'ire_legend.png'")
|
|
|
|
return results
|
|
|
|
def test_shadow_transition_analysis():
|
|
"""
|
|
Test shadow transition analysis specifically.
|
|
"""
|
|
print("\n=== Shadow Transition Analysis Test ===")
|
|
|
|
# Create test image with known shadow characteristics
|
|
test_image = create_test_image()
|
|
|
|
analyzer = IREShadowAnalyzer(shadow_ire_threshold=20.0, transition_sensitivity=0.1)
|
|
|
|
# Convert to IRE
|
|
ire_values = analyzer.srgb_to_ire(test_image)
|
|
|
|
# Analyze transitions
|
|
transition_analysis = analyzer.analyze_shadow_transitions(ire_values)
|
|
|
|
print("Transition Analysis Results:")
|
|
print(f"Mean Shadow Gradient: {transition_analysis['mean_shadow_gradient']:.2f}")
|
|
print(f"Max Shadow Gradient: {transition_analysis['max_shadow_gradient']:.2f}")
|
|
print(f"Soft Shadow Ratio: {transition_analysis['soft_shadow_ratio']:.3f}")
|
|
print(f"Hard Shadow Ratio: {transition_analysis['hard_shadow_ratio']:.3f}")
|
|
|
|
# Classify shadow characteristics
|
|
classification = analyzer.classify_shadow_characteristics(transition_analysis)
|
|
|
|
print(f"\nShadow Classification:")
|
|
print(f"Shadow Character: {classification['shadow_character']}")
|
|
print(f"Transition Quality: {classification['transition_quality']}")
|
|
|
|
return transition_analysis, classification
|
|
|
|
def demonstrate_ire_color_mapping():
|
|
"""
|
|
Demonstrate the IRE color mapping system.
|
|
"""
|
|
print("\n=== IRE Color Mapping Demonstration ===")
|
|
|
|
analyzer = IREShadowAnalyzer()
|
|
|
|
# Create a test image with known IRE values
|
|
test_ire_values = torch.linspace(0, 100, 100).unsqueeze(0).unsqueeze(0).repeat(1, 50, 1)
|
|
test_ire_values = test_ire_values.unsqueeze(0) # Add batch dimension
|
|
|
|
# Create false color visualization
|
|
false_color = analyzer.create_false_color_visualization(test_ire_values)
|
|
|
|
# Convert to numpy for visualization
|
|
false_color_np = false_color[0].numpy()
|
|
|
|
plt.figure(figsize=(12, 4))
|
|
plt.imshow(false_color_np)
|
|
plt.title('IRE Color Mapping Demonstration (0-100 IRE)')
|
|
plt.xlabel('IRE Value')
|
|
plt.ylabel('')
|
|
plt.xticks(range(0, 100, 10), range(0, 100, 10))
|
|
plt.yticks([])
|
|
|
|
# Add color zone labels
|
|
zone_boundaries = [0, 10, 20, 30, 50, 70, 90, 100]
|
|
zone_labels = ['Deep Shadow', 'Shadow', 'Mid-Shadow', 'Low-Midtone',
|
|
'Midtone', 'High-Midtone', 'Highlight', 'Bright Highlight']
|
|
|
|
for i, (boundary, label) in enumerate(zip(zone_boundaries[:-1], zone_labels)):
|
|
plt.axvline(x=boundary, color='white', linestyle='--', alpha=0.7)
|
|
plt.text(boundary + 5, 25, label, rotation=90, verticalalignment='bottom',
|
|
fontsize=8, color='white', weight='bold')
|
|
|
|
plt.tight_layout()
|
|
plt.savefig('ire_color_mapping_demo.png', dpi=150, bbox_inches='tight')
|
|
print("IRE color mapping demonstration saved to 'ire_color_mapping_demo.png'")
|
|
|
|
if __name__ == "__main__":
|
|
print("Starting IRE Shadow Analysis Tests...")
|
|
|
|
try:
|
|
# Test basic IRE analysis
|
|
results = test_ire_analysis()
|
|
|
|
# Test shadow transition analysis
|
|
transition_results, classification = test_shadow_transition_analysis()
|
|
|
|
# Demonstrate IRE color mapping
|
|
demonstrate_ire_color_mapping()
|
|
|
|
print("\n=== All Tests Completed Successfully ===")
|
|
print("Generated files:")
|
|
print("- ire_analysis_results.png: Complete analysis visualization")
|
|
print("- ire_legend.png: IRE color mapping legend")
|
|
print("- ire_color_mapping_demo.png: IRE color mapping demonstration")
|
|
|
|
except Exception as e:
|
|
print(f"Error during testing: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|