Implement WiFi-DensePose system with CSI data extraction and router interface
- Added CSIExtractor class for extracting CSI data from WiFi routers. - Implemented RouterInterface class for SSH communication with routers. - Developed DensePoseHead class for body part segmentation and UV coordinate regression. - Created unit tests for CSIExtractor and RouterInterface to ensure functionality and error handling. - Integrated paramiko for SSH connections and command execution. - Established configuration validation for both extractor and router interface. - Added context manager support for resource management in both classes.
This commit is contained in:
rUv
@@ -3,27 +3,27 @@ import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
from unittest.mock import Mock, patch
|
||||
from src.models.densepose_head import DensePoseHead
|
||||
from src.models.densepose_head import DensePoseHead, DensePoseError
|
||||
|
||||
|
||||
class TestDensePoseHead:
|
||||
"""Test suite for DensePose Head following London School TDD principles"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_feature_input(self):
|
||||
"""Generate synthetic feature input tensor for testing"""
|
||||
# Batch size 2, 512 channels, 56 height, 100 width (from modality translation)
|
||||
return torch.randn(2, 512, 56, 100)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_config(self):
|
||||
"""Configuration for DensePose head"""
|
||||
return {
|
||||
'input_channels': 512,
|
||||
'num_body_parts': 24, # Standard DensePose body parts
|
||||
'num_uv_coordinates': 2, # U and V coordinates
|
||||
'hidden_dim': 256,
|
||||
'dropout_rate': 0.1
|
||||
'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
|
||||
@@ -31,6 +31,33 @@ class TestDensePoseHead:
|
||||
"""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
|
||||
@@ -39,135 +66,302 @@ class TestDensePoseHead:
|
||||
# Assert
|
||||
assert head is not None
|
||||
assert isinstance(head, nn.Module)
|
||||
assert hasattr(head, 'segmentation_head')
|
||||
assert hasattr(head, 'uv_regression_head')
|
||||
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_shapes(self, densepose_head, mock_feature_input):
|
||||
"""Test that forward pass produces correctly shaped outputs"""
|
||||
def test_forward_pass_produces_correct_output_format(self, densepose_head, mock_feature_input):
|
||||
"""Test that forward pass produces correctly formatted output"""
|
||||
# Act
|
||||
with torch.no_grad():
|
||||
segmentation, uv_coords = densepose_head(mock_feature_input)
|
||||
output = densepose_head(mock_feature_input)
|
||||
|
||||
# Assert
|
||||
assert segmentation is not None
|
||||
assert uv_coords is not None
|
||||
assert isinstance(segmentation, torch.Tensor)
|
||||
assert isinstance(uv_coords, torch.Tensor)
|
||||
assert output is not None
|
||||
assert isinstance(output, dict)
|
||||
assert 'segmentation' in output
|
||||
assert 'uv_coordinates' in output
|
||||
|
||||
# Check segmentation output shape
|
||||
assert segmentation.shape[0] == mock_feature_input.shape[0] # Batch size preserved
|
||||
assert segmentation.shape[1] == densepose_head.num_body_parts # Correct number of body parts
|
||||
assert segmentation.shape[2:] == mock_feature_input.shape[2:] # Spatial dimensions preserved
|
||||
seg_output = output['segmentation']
|
||||
uv_output = output['uv_coordinates']
|
||||
|
||||
# Check UV coordinates output shape
|
||||
assert uv_coords.shape[0] == mock_feature_input.shape[0] # Batch size preserved
|
||||
assert uv_coords.shape[1] == densepose_head.num_uv_coordinates # U and V coordinates
|
||||
assert uv_coords.shape[2:] == mock_feature_input.shape[2:] # Spatial dimensions preserved
|
||||
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_output_has_valid_probabilities(self, densepose_head, mock_feature_input):
|
||||
"""Test that segmentation output has valid probability distributions"""
|
||||
def test_segmentation_head_produces_correct_shape(self, densepose_head, mock_feature_input):
|
||||
"""Test that segmentation head produces correct output shape"""
|
||||
# Act
|
||||
with torch.no_grad():
|
||||
segmentation, _ = densepose_head(mock_feature_input)
|
||||
output = densepose_head(mock_feature_input)
|
||||
seg_output = output['segmentation']
|
||||
|
||||
# Assert
|
||||
# After softmax, values should be between 0 and 1
|
||||
assert torch.all(segmentation >= 0.0)
|
||||
assert torch.all(segmentation <= 1.0)
|
||||
|
||||
# Sum across body parts dimension should be approximately 1
|
||||
part_sums = torch.sum(segmentation, dim=1)
|
||||
assert torch.allclose(part_sums, torch.ones_like(part_sums), atol=1e-5)
|
||||
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_coordinates_output_in_valid_range(self, densepose_head, mock_feature_input):
|
||||
"""Test that UV coordinates are in valid range [0, 1]"""
|
||||
def test_uv_regression_head_produces_correct_shape(self, densepose_head, mock_feature_input):
|
||||
"""Test that UV regression head produces correct output shape"""
|
||||
# Act
|
||||
with torch.no_grad():
|
||||
_, uv_coords = densepose_head(mock_feature_input)
|
||||
output = densepose_head(mock_feature_input)
|
||||
uv_output = output['uv_coordinates']
|
||||
|
||||
# Assert
|
||||
# UV coordinates should be in range [0, 1] after sigmoid
|
||||
assert torch.all(uv_coords >= 0.0)
|
||||
assert torch.all(uv_coords <= 1.0)
|
||||
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_head_handles_different_batch_sizes(self, densepose_head):
|
||||
"""Test that head handles different batch sizes correctly"""
|
||||
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
|
||||
batch_sizes = [1, 4, 8]
|
||||
output = densepose_head(mock_feature_input)
|
||||
seg_logits = output['segmentation']
|
||||
|
||||
for batch_size in batch_sizes:
|
||||
input_tensor = torch.randn(batch_size, 512, 56, 100)
|
||||
|
||||
# Act
|
||||
with torch.no_grad():
|
||||
segmentation, uv_coords = densepose_head(input_tensor)
|
||||
|
||||
# Assert
|
||||
assert segmentation.shape[0] == batch_size
|
||||
assert uv_coords.shape[0] == batch_size
|
||||
|
||||
def test_head_is_trainable(self, densepose_head, mock_feature_input):
|
||||
"""Test that head parameters are trainable"""
|
||||
# Arrange
|
||||
seg_criterion = nn.CrossEntropyLoss()
|
||||
uv_criterion = nn.MSELoss()
|
||||
|
||||
# Create targets with correct shapes
|
||||
seg_target = torch.randint(0, 24, (2, 56, 100)) # Class indices for segmentation
|
||||
uv_target = torch.rand(2, 2, 56, 100) # UV coordinates target
|
||||
# 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
|
||||
segmentation, uv_coords = densepose_head(mock_feature_input)
|
||||
seg_loss = seg_criterion(segmentation, seg_target)
|
||||
uv_loss = uv_criterion(uv_coords, uv_target)
|
||||
total_loss = seg_loss + uv_loss
|
||||
total_loss.backward()
|
||||
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
|
||||
# Check that gradients are computed
|
||||
# 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_handles_invalid_input_shape(self, densepose_head):
|
||||
"""Test that head handles invalid input shapes gracefully"""
|
||||
def test_head_configuration_validation(self):
|
||||
"""Test that head validates configuration parameters"""
|
||||
# Arrange
|
||||
invalid_input = torch.randn(2, 256, 56, 100) # Wrong number of channels
|
||||
invalid_config = {
|
||||
'input_channels': 0, # Invalid
|
||||
'num_body_parts': -1, # Invalid
|
||||
'num_uv_coordinates': 2
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(RuntimeError):
|
||||
densepose_head(invalid_input)
|
||||
with pytest.raises(ValueError):
|
||||
DensePoseHead(invalid_config)
|
||||
|
||||
def test_head_supports_evaluation_mode(self, densepose_head, mock_feature_input):
|
||||
"""Test that head supports evaluation mode"""
|
||||
# Act
|
||||
densepose_head.eval()
|
||||
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)
|
||||
|
||||
with torch.no_grad():
|
||||
seg1, uv1 = densepose_head(mock_feature_input)
|
||||
seg2, uv2 = densepose_head(mock_feature_input)
|
||||
# Act - Save state
|
||||
state_dict = densepose_head.state_dict()
|
||||
|
||||
# Assert - In eval mode with same input, outputs should be identical
|
||||
assert torch.allclose(seg1, seg2, atol=1e-6)
|
||||
assert torch.allclose(uv1, uv2, atol=1e-6)
|
||||
|
||||
def test_head_output_quality(self, densepose_head, mock_feature_input):
|
||||
"""Test that head produces meaningful outputs"""
|
||||
# Act
|
||||
with torch.no_grad():
|
||||
segmentation, uv_coords = densepose_head(mock_feature_input)
|
||||
# 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
|
||||
# Outputs should not contain NaN or Inf values
|
||||
assert not torch.isnan(segmentation).any()
|
||||
assert not torch.isinf(segmentation).any()
|
||||
assert not torch.isnan(uv_coords).any()
|
||||
assert not torch.isinf(uv_coords).any()
|
||||
|
||||
# Outputs should have reasonable variance (not all zeros or ones)
|
||||
assert segmentation.std() > 0.01
|
||||
assert uv_coords.std() > 0.01
|
||||
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)
|
||||
Reference in New Issue
Block a user