import pytest import torch import torch.nn as nn import numpy as np from unittest.mock import Mock, patch from src.models.densepose_head import DensePoseHead, DensePoseError class TestDensePoseHead: """Test suite for DensePose Head following London School TDD principles""" @pytest.fixture def mock_config(self): """Configuration for DensePose head""" return { 'input_channels': 256, 'num_body_parts': 24, 'num_uv_coordinates': 2, 'hidden_channels': [128, 64], 'kernel_size': 3, 'padding': 1, 'dropout_rate': 0.1, 'use_deformable_conv': False, 'use_fpn': True, 'fpn_levels': [2, 3, 4, 5], 'output_stride': 4 } @pytest.fixture def densepose_head(self, mock_config): """Create DensePose head instance for testing""" return DensePoseHead(mock_config) @pytest.fixture def mock_feature_input(self): """Generate mock feature input tensor""" batch_size = 2 channels = 256 height = 56 width = 56 return torch.randn(batch_size, channels, height, width) @pytest.fixture def mock_target_masks(self): """Generate mock target segmentation masks""" batch_size = 2 num_parts = 24 height = 224 width = 224 return torch.randint(0, num_parts + 1, (batch_size, height, width)) @pytest.fixture def mock_target_uv(self): """Generate mock target UV coordinates""" batch_size = 2 num_coords = 2 height = 224 width = 224 return torch.randn(batch_size, num_coords, height, width) def test_head_initialization_creates_correct_architecture(self, mock_config): """Test that DensePose head initializes with correct architecture""" # Act head = DensePoseHead(mock_config) # Assert assert head is not None assert isinstance(head, nn.Module) assert head.input_channels == mock_config['input_channels'] assert head.num_body_parts == mock_config['num_body_parts'] assert head.num_uv_coordinates == mock_config['num_uv_coordinates'] assert head.use_fpn == mock_config['use_fpn'] assert hasattr(head, 'segmentation_head') assert hasattr(head, 'uv_regression_head') if mock_config['use_fpn']: assert hasattr(head, 'fpn') def test_forward_pass_produces_correct_output_format(self, densepose_head, mock_feature_input): """Test that forward pass produces correctly formatted output""" # Act output = densepose_head(mock_feature_input) # Assert assert output is not None assert isinstance(output, dict) assert 'segmentation' in output assert 'uv_coordinates' in output seg_output = output['segmentation'] uv_output = output['uv_coordinates'] assert isinstance(seg_output, torch.Tensor) assert isinstance(uv_output, torch.Tensor) assert seg_output.shape[0] == mock_feature_input.shape[0] # Batch size preserved assert uv_output.shape[0] == mock_feature_input.shape[0] # Batch size preserved def test_segmentation_head_produces_correct_shape(self, densepose_head, mock_feature_input): """Test that segmentation head produces correct output shape""" # Act output = densepose_head(mock_feature_input) seg_output = output['segmentation'] # Assert expected_channels = densepose_head.num_body_parts + 1 # +1 for background assert seg_output.shape[1] == expected_channels assert seg_output.shape[2] >= mock_feature_input.shape[2] # Height upsampled assert seg_output.shape[3] >= mock_feature_input.shape[3] # Width upsampled def test_uv_regression_head_produces_correct_shape(self, densepose_head, mock_feature_input): """Test that UV regression head produces correct output shape""" # Act output = densepose_head(mock_feature_input) uv_output = output['uv_coordinates'] # Assert assert uv_output.shape[1] == densepose_head.num_uv_coordinates assert uv_output.shape[2] >= mock_feature_input.shape[2] # Height upsampled assert uv_output.shape[3] >= mock_feature_input.shape[3] # Width upsampled def test_compute_segmentation_loss_measures_pixel_classification(self, densepose_head, mock_feature_input, mock_target_masks): """Test that compute_segmentation_loss measures pixel classification accuracy""" # Arrange output = densepose_head(mock_feature_input) seg_logits = output['segmentation'] # Resize target to match output target_resized = torch.nn.functional.interpolate( mock_target_masks.float().unsqueeze(1), size=seg_logits.shape[2:], mode='nearest' ).squeeze(1).long() # Act loss = densepose_head.compute_segmentation_loss(seg_logits, target_resized) # Assert assert loss is not None assert isinstance(loss, torch.Tensor) assert loss.dim() == 0 # Scalar loss assert loss.item() >= 0 # Loss should be non-negative def test_compute_uv_loss_measures_coordinate_regression(self, densepose_head, mock_feature_input, mock_target_uv): """Test that compute_uv_loss measures UV coordinate regression accuracy""" # Arrange output = densepose_head(mock_feature_input) uv_pred = output['uv_coordinates'] # Resize target to match output target_resized = torch.nn.functional.interpolate( mock_target_uv, size=uv_pred.shape[2:], mode='bilinear', align_corners=False ) # Act loss = densepose_head.compute_uv_loss(uv_pred, target_resized) # Assert assert loss is not None assert isinstance(loss, torch.Tensor) assert loss.dim() == 0 # Scalar loss assert loss.item() >= 0 # Loss should be non-negative def test_compute_total_loss_combines_segmentation_and_uv_losses(self, densepose_head, mock_feature_input, mock_target_masks, mock_target_uv): """Test that compute_total_loss combines segmentation and UV losses""" # Arrange output = densepose_head(mock_feature_input) # Resize targets to match outputs seg_target = torch.nn.functional.interpolate( mock_target_masks.float().unsqueeze(1), size=output['segmentation'].shape[2:], mode='nearest' ).squeeze(1).long() uv_target = torch.nn.functional.interpolate( mock_target_uv, size=output['uv_coordinates'].shape[2:], mode='bilinear', align_corners=False ) # Act total_loss = densepose_head.compute_total_loss(output, seg_target, uv_target) seg_loss = densepose_head.compute_segmentation_loss(output['segmentation'], seg_target) uv_loss = densepose_head.compute_uv_loss(output['uv_coordinates'], uv_target) # Assert assert total_loss is not None assert isinstance(total_loss, torch.Tensor) assert total_loss.item() > 0 # Total loss should be combination of individual losses expected_total = seg_loss + uv_loss assert torch.allclose(total_loss, expected_total, atol=1e-6) def test_fpn_integration_enhances_multi_scale_features(self, mock_config, mock_feature_input): """Test that FPN integration enhances multi-scale feature processing""" # Arrange config_with_fpn = mock_config.copy() config_with_fpn['use_fpn'] = True config_without_fpn = mock_config.copy() config_without_fpn['use_fpn'] = False head_with_fpn = DensePoseHead(config_with_fpn) head_without_fpn = DensePoseHead(config_without_fpn) # Act output_with_fpn = head_with_fpn(mock_feature_input) output_without_fpn = head_without_fpn(mock_feature_input) # Assert assert output_with_fpn['segmentation'].shape == output_without_fpn['segmentation'].shape assert output_with_fpn['uv_coordinates'].shape == output_without_fpn['uv_coordinates'].shape # Outputs should be different due to FPN assert not torch.allclose(output_with_fpn['segmentation'], output_without_fpn['segmentation'], atol=1e-6) def test_get_prediction_confidence_provides_uncertainty_estimates(self, densepose_head, mock_feature_input): """Test that get_prediction_confidence provides uncertainty estimates""" # Arrange output = densepose_head(mock_feature_input) # Act confidence = densepose_head.get_prediction_confidence(output) # Assert assert confidence is not None assert isinstance(confidence, dict) assert 'segmentation_confidence' in confidence assert 'uv_confidence' in confidence seg_conf = confidence['segmentation_confidence'] uv_conf = confidence['uv_confidence'] assert isinstance(seg_conf, torch.Tensor) assert isinstance(uv_conf, torch.Tensor) assert seg_conf.shape[0] == mock_feature_input.shape[0] assert uv_conf.shape[0] == mock_feature_input.shape[0] def test_post_process_predictions_formats_output(self, densepose_head, mock_feature_input): """Test that post_process_predictions formats output correctly""" # Arrange raw_output = densepose_head(mock_feature_input) # Act processed = densepose_head.post_process_predictions(raw_output) # Assert assert processed is not None assert isinstance(processed, dict) assert 'body_parts' in processed assert 'uv_coordinates' in processed assert 'confidence_scores' in processed def test_training_mode_enables_dropout(self, densepose_head, mock_feature_input): """Test that training mode enables dropout for regularization""" # Arrange densepose_head.train() # Act output1 = densepose_head(mock_feature_input) output2 = densepose_head(mock_feature_input) # Assert - outputs should be different due to dropout assert not torch.allclose(output1['segmentation'], output2['segmentation'], atol=1e-6) assert not torch.allclose(output1['uv_coordinates'], output2['uv_coordinates'], atol=1e-6) def test_evaluation_mode_disables_dropout(self, densepose_head, mock_feature_input): """Test that evaluation mode disables dropout for consistent inference""" # Arrange densepose_head.eval() # Act output1 = densepose_head(mock_feature_input) output2 = densepose_head(mock_feature_input) # Assert - outputs should be identical in eval mode assert torch.allclose(output1['segmentation'], output2['segmentation'], atol=1e-6) assert torch.allclose(output1['uv_coordinates'], output2['uv_coordinates'], atol=1e-6) def test_head_validates_input_dimensions(self, densepose_head): """Test that head validates input dimensions""" # Arrange invalid_input = torch.randn(2, 128, 56, 56) # Wrong number of channels # Act & Assert with pytest.raises(DensePoseError): densepose_head(invalid_input) def test_head_handles_different_input_sizes(self, densepose_head): """Test that head handles different input sizes""" # Arrange small_input = torch.randn(1, 256, 28, 28) large_input = torch.randn(1, 256, 112, 112) # Act small_output = densepose_head(small_input) large_output = densepose_head(large_input) # Assert assert small_output['segmentation'].shape[2:] != large_output['segmentation'].shape[2:] assert small_output['uv_coordinates'].shape[2:] != large_output['uv_coordinates'].shape[2:] def test_head_supports_gradient_computation(self, densepose_head, mock_feature_input, mock_target_masks, mock_target_uv): """Test that head supports gradient computation for training""" # Arrange densepose_head.train() optimizer = torch.optim.Adam(densepose_head.parameters(), lr=0.001) output = densepose_head(mock_feature_input) # Resize targets seg_target = torch.nn.functional.interpolate( mock_target_masks.float().unsqueeze(1), size=output['segmentation'].shape[2:], mode='nearest' ).squeeze(1).long() uv_target = torch.nn.functional.interpolate( mock_target_uv, size=output['uv_coordinates'].shape[2:], mode='bilinear', align_corners=False ) # Act loss = densepose_head.compute_total_loss(output, seg_target, uv_target) optimizer.zero_grad() loss.backward() # Assert for param in densepose_head.parameters(): if param.requires_grad: assert param.grad is not None assert not torch.allclose(param.grad, torch.zeros_like(param.grad)) def test_head_configuration_validation(self): """Test that head validates configuration parameters""" # Arrange invalid_config = { 'input_channels': 0, # Invalid 'num_body_parts': -1, # Invalid 'num_uv_coordinates': 2 } # Act & Assert with pytest.raises(ValueError): DensePoseHead(invalid_config) def test_save_and_load_model_state(self, densepose_head, mock_feature_input): """Test that model state can be saved and loaded""" # Arrange original_output = densepose_head(mock_feature_input) # Act - Save state state_dict = densepose_head.state_dict() # Create new head and load state new_head = DensePoseHead(densepose_head.config) new_head.load_state_dict(state_dict) new_output = new_head(mock_feature_input) # Assert assert torch.allclose(original_output['segmentation'], new_output['segmentation'], atol=1e-6) assert torch.allclose(original_output['uv_coordinates'], new_output['uv_coordinates'], atol=1e-6)