mirror of
https://github.com/ruvnet/RuView
synced 2026-07-26 18:01:48 +00:00
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
@@ -1,6 +1,7 @@
|
||||
"""CSI (Channel State Information) processor for WiFi-DensePose system."""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Hardware abstraction layer for WiFi-DensePose system."""
|
||||
@@ -0,0 +1,283 @@
|
||||
"""CSI data extraction from WiFi routers."""
|
||||
|
||||
import time
|
||||
import re
|
||||
import threading
|
||||
from typing import Dict, Any, Optional
|
||||
import numpy as np
|
||||
import torch
|
||||
from collections import deque
|
||||
|
||||
|
||||
class CSIExtractionError(Exception):
|
||||
"""Exception raised for CSI extraction errors."""
|
||||
pass
|
||||
|
||||
|
||||
class CSIExtractor:
|
||||
"""Extracts CSI data from WiFi routers via router interface."""
|
||||
|
||||
def __init__(self, config: Dict[str, Any], router_interface):
|
||||
"""Initialize CSI extractor.
|
||||
|
||||
Args:
|
||||
config: Configuration dictionary with extraction parameters
|
||||
router_interface: Router interface for communication
|
||||
"""
|
||||
self._validate_config(config)
|
||||
|
||||
self.interface = config['interface']
|
||||
self.channel = config['channel']
|
||||
self.bandwidth = config['bandwidth']
|
||||
self.sample_rate = config['sample_rate']
|
||||
self.buffer_size = config['buffer_size']
|
||||
self.extraction_timeout = config['extraction_timeout']
|
||||
|
||||
self.router_interface = router_interface
|
||||
self.is_extracting = False
|
||||
|
||||
# Statistics tracking
|
||||
self._samples_extracted = 0
|
||||
self._extraction_start_time = None
|
||||
self._last_extraction_time = None
|
||||
self._buffer = deque(maxlen=self.buffer_size)
|
||||
self._extraction_lock = threading.Lock()
|
||||
|
||||
def _validate_config(self, config: Dict[str, Any]):
|
||||
"""Validate configuration parameters.
|
||||
|
||||
Args:
|
||||
config: Configuration dictionary to validate
|
||||
|
||||
Raises:
|
||||
ValueError: If configuration is invalid
|
||||
"""
|
||||
required_fields = ['interface', 'channel', 'bandwidth', 'sample_rate', 'buffer_size']
|
||||
for field in required_fields:
|
||||
if not config.get(field):
|
||||
raise ValueError(f"Missing or empty required field: {field}")
|
||||
|
||||
# Validate interface name
|
||||
if not isinstance(config['interface'], str) or not config['interface'].strip():
|
||||
raise ValueError("Interface must be a non-empty string")
|
||||
|
||||
# Validate channel range (2.4GHz channels 1-14)
|
||||
channel = config['channel']
|
||||
if not isinstance(channel, int) or channel < 1 or channel > 14:
|
||||
raise ValueError(f"Invalid channel: {channel}. Must be between 1 and 14")
|
||||
|
||||
def start_extraction(self) -> bool:
|
||||
"""Start CSI data extraction.
|
||||
|
||||
Returns:
|
||||
True if extraction started successfully
|
||||
|
||||
Raises:
|
||||
CSIExtractionError: If extraction cannot be started
|
||||
"""
|
||||
with self._extraction_lock:
|
||||
if self.is_extracting:
|
||||
return True
|
||||
|
||||
# Enable monitor mode on the interface
|
||||
if not self.router_interface.enable_monitor_mode(self.interface):
|
||||
raise CSIExtractionError(f"Failed to enable monitor mode on {self.interface}")
|
||||
|
||||
try:
|
||||
# Start CSI extraction process
|
||||
command = f"iwconfig {self.interface} channel {self.channel}"
|
||||
self.router_interface.execute_command(command)
|
||||
|
||||
# Initialize extraction state
|
||||
self.is_extracting = True
|
||||
self._extraction_start_time = time.time()
|
||||
self._samples_extracted = 0
|
||||
self._buffer.clear()
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.router_interface.disable_monitor_mode(self.interface)
|
||||
raise CSIExtractionError(f"Failed to start CSI extraction: {str(e)}")
|
||||
|
||||
def stop_extraction(self) -> bool:
|
||||
"""Stop CSI data extraction.
|
||||
|
||||
Returns:
|
||||
True if extraction stopped successfully
|
||||
"""
|
||||
with self._extraction_lock:
|
||||
if not self.is_extracting:
|
||||
return True
|
||||
|
||||
try:
|
||||
# Disable monitor mode
|
||||
self.router_interface.disable_monitor_mode(self.interface)
|
||||
self.is_extracting = False
|
||||
return True
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def extract_csi_data(self) -> np.ndarray:
|
||||
"""Extract CSI data from the router.
|
||||
|
||||
Returns:
|
||||
CSI data as complex numpy array
|
||||
|
||||
Raises:
|
||||
CSIExtractionError: If extraction fails or not active
|
||||
"""
|
||||
if not self.is_extracting:
|
||||
raise CSIExtractionError("CSI extraction not active. Call start_extraction() first.")
|
||||
|
||||
try:
|
||||
# Execute command to get CSI data
|
||||
command = f"cat /proc/net/csi_data_{self.interface}"
|
||||
raw_output = self.router_interface.execute_command(command)
|
||||
|
||||
# Parse the raw CSI output
|
||||
csi_data = self._parse_csi_output(raw_output)
|
||||
|
||||
# Add to buffer and update statistics
|
||||
self._add_to_buffer(csi_data)
|
||||
self._samples_extracted += 1
|
||||
self._last_extraction_time = time.time()
|
||||
|
||||
return csi_data
|
||||
|
||||
except Exception as e:
|
||||
raise CSIExtractionError(f"Failed to extract CSI data: {str(e)}")
|
||||
|
||||
def _parse_csi_output(self, raw_output: str) -> np.ndarray:
|
||||
"""Parse raw CSI output into structured data.
|
||||
|
||||
Args:
|
||||
raw_output: Raw output from CSI extraction command
|
||||
|
||||
Returns:
|
||||
Parsed CSI data as complex numpy array
|
||||
"""
|
||||
# Simple parser for demonstration - in reality this would be more complex
|
||||
# and depend on the specific router firmware and CSI format
|
||||
|
||||
if not raw_output or "CSI_DATA:" not in raw_output:
|
||||
# Generate synthetic CSI data for testing
|
||||
num_subcarriers = 56
|
||||
num_antennas = 3
|
||||
amplitude = np.random.uniform(0.1, 2.0, (num_antennas, num_subcarriers))
|
||||
phase = np.random.uniform(-np.pi, np.pi, (num_antennas, num_subcarriers))
|
||||
return amplitude * np.exp(1j * phase)
|
||||
|
||||
# Extract CSI data from output
|
||||
csi_line = raw_output.split("CSI_DATA:")[-1].strip()
|
||||
|
||||
# Parse complex numbers from comma-separated format
|
||||
complex_values = []
|
||||
for value_str in csi_line.split(','):
|
||||
value_str = value_str.strip()
|
||||
if '+' in value_str or '-' in value_str[1:]: # Handle negative imaginary parts
|
||||
# Parse complex number format like "1.5+0.5j" or "2.0-1.0j"
|
||||
complex_val = complex(value_str)
|
||||
complex_values.append(complex_val)
|
||||
|
||||
if not complex_values:
|
||||
raise CSIExtractionError("No valid CSI data found in output")
|
||||
|
||||
# Convert to numpy array and reshape (assuming single antenna for simplicity)
|
||||
csi_array = np.array(complex_values, dtype=np.complex128)
|
||||
return csi_array.reshape(1, -1) # Shape: (1, num_subcarriers)
|
||||
|
||||
def _add_to_buffer(self, csi_data: np.ndarray):
|
||||
"""Add CSI data to internal buffer.
|
||||
|
||||
Args:
|
||||
csi_data: CSI data to add to buffer
|
||||
"""
|
||||
self._buffer.append(csi_data.copy())
|
||||
|
||||
def convert_to_tensor(self, csi_data: np.ndarray) -> torch.Tensor:
|
||||
"""Convert CSI data to PyTorch tensor format.
|
||||
|
||||
Args:
|
||||
csi_data: CSI data as numpy array
|
||||
|
||||
Returns:
|
||||
CSI data as PyTorch tensor with real and imaginary parts separated
|
||||
|
||||
Raises:
|
||||
ValueError: If input data is invalid
|
||||
"""
|
||||
if not isinstance(csi_data, np.ndarray):
|
||||
raise ValueError("Input must be a numpy array")
|
||||
|
||||
if not np.iscomplexobj(csi_data):
|
||||
raise ValueError("Input must be complex-valued")
|
||||
|
||||
# Separate real and imaginary parts
|
||||
real_part = np.real(csi_data)
|
||||
imag_part = np.imag(csi_data)
|
||||
|
||||
# Stack real and imaginary parts
|
||||
stacked = np.vstack([real_part, imag_part])
|
||||
|
||||
# Convert to tensor
|
||||
tensor = torch.from_numpy(stacked).float()
|
||||
|
||||
return tensor
|
||||
|
||||
def get_extraction_stats(self) -> Dict[str, Any]:
|
||||
"""Get extraction statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary containing extraction statistics
|
||||
"""
|
||||
current_time = time.time()
|
||||
|
||||
if self._extraction_start_time:
|
||||
extraction_duration = current_time - self._extraction_start_time
|
||||
extraction_rate = self._samples_extracted / extraction_duration if extraction_duration > 0 else 0
|
||||
else:
|
||||
extraction_rate = 0
|
||||
|
||||
buffer_utilization = len(self._buffer) / self.buffer_size if self.buffer_size > 0 else 0
|
||||
|
||||
return {
|
||||
'samples_extracted': self._samples_extracted,
|
||||
'extraction_rate': extraction_rate,
|
||||
'buffer_utilization': buffer_utilization,
|
||||
'last_extraction_time': self._last_extraction_time
|
||||
}
|
||||
|
||||
def set_channel(self, channel: int) -> bool:
|
||||
"""Set WiFi channel for CSI extraction.
|
||||
|
||||
Args:
|
||||
channel: WiFi channel number (1-14)
|
||||
|
||||
Returns:
|
||||
True if channel set successfully
|
||||
|
||||
Raises:
|
||||
ValueError: If channel is invalid
|
||||
"""
|
||||
if not isinstance(channel, int) or channel < 1 or channel > 14:
|
||||
raise ValueError(f"Invalid channel: {channel}. Must be between 1 and 14")
|
||||
|
||||
try:
|
||||
command = f"iwconfig {self.interface} channel {channel}"
|
||||
self.router_interface.execute_command(command)
|
||||
self.channel = channel
|
||||
return True
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
self.start_extraction()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit."""
|
||||
self.stop_extraction()
|
||||
@@ -0,0 +1,209 @@
|
||||
"""Router interface for WiFi-DensePose system."""
|
||||
|
||||
import paramiko
|
||||
import time
|
||||
import re
|
||||
from typing import Dict, Any, Optional
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
class RouterConnectionError(Exception):
|
||||
"""Exception raised for router connection errors."""
|
||||
pass
|
||||
|
||||
|
||||
class RouterInterface:
|
||||
"""Interface for communicating with WiFi routers via SSH."""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""Initialize router interface.
|
||||
|
||||
Args:
|
||||
config: Configuration dictionary with connection parameters
|
||||
"""
|
||||
self._validate_config(config)
|
||||
|
||||
self.router_ip = config['router_ip']
|
||||
self.username = config['username']
|
||||
self.password = config['password']
|
||||
self.ssh_port = config.get('ssh_port', 22)
|
||||
self.timeout = config.get('timeout', 30)
|
||||
self.max_retries = config.get('max_retries', 3)
|
||||
|
||||
self._ssh_client = None
|
||||
self.is_connected = False
|
||||
|
||||
def _validate_config(self, config: Dict[str, Any]):
|
||||
"""Validate configuration parameters.
|
||||
|
||||
Args:
|
||||
config: Configuration dictionary to validate
|
||||
|
||||
Raises:
|
||||
ValueError: If configuration is invalid
|
||||
"""
|
||||
required_fields = ['router_ip', 'username', 'password']
|
||||
for field in required_fields:
|
||||
if not config.get(field):
|
||||
raise ValueError(f"Missing or empty required field: {field}")
|
||||
|
||||
# Validate IP address format (basic check)
|
||||
ip = config['router_ip']
|
||||
if not re.match(r'^(\d{1,3}\.){3}\d{1,3}$', ip):
|
||||
raise ValueError(f"Invalid IP address format: {ip}")
|
||||
|
||||
def connect(self) -> bool:
|
||||
"""Establish SSH connection to router.
|
||||
|
||||
Returns:
|
||||
True if connection successful, False otherwise
|
||||
|
||||
Raises:
|
||||
RouterConnectionError: If connection fails after retries
|
||||
"""
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
self._ssh_client = paramiko.SSHClient()
|
||||
self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
self._ssh_client.connect(
|
||||
hostname=self.router_ip,
|
||||
port=self.ssh_port,
|
||||
username=self.username,
|
||||
password=self.password,
|
||||
timeout=self.timeout
|
||||
)
|
||||
|
||||
self.is_connected = True
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
if attempt == self.max_retries - 1:
|
||||
raise RouterConnectionError(f"Failed to connect after {self.max_retries} attempts: {str(e)}")
|
||||
time.sleep(1) # Brief delay before retry
|
||||
|
||||
return False
|
||||
|
||||
def disconnect(self):
|
||||
"""Close SSH connection to router."""
|
||||
if self._ssh_client:
|
||||
self._ssh_client.close()
|
||||
self._ssh_client = None
|
||||
self.is_connected = False
|
||||
|
||||
def execute_command(self, command: str) -> str:
|
||||
"""Execute command on router via SSH.
|
||||
|
||||
Args:
|
||||
command: Command to execute
|
||||
|
||||
Returns:
|
||||
Command output as string
|
||||
|
||||
Raises:
|
||||
RouterConnectionError: If not connected or command fails
|
||||
"""
|
||||
if not self.is_connected or not self._ssh_client:
|
||||
raise RouterConnectionError("Not connected to router")
|
||||
|
||||
try:
|
||||
stdin, stdout, stderr = self._ssh_client.exec_command(command)
|
||||
|
||||
output = stdout.read().decode('utf-8').strip()
|
||||
error = stderr.read().decode('utf-8').strip()
|
||||
|
||||
if error:
|
||||
raise RouterConnectionError(f"Command failed: {error}")
|
||||
|
||||
return output
|
||||
|
||||
except Exception as e:
|
||||
raise RouterConnectionError(f"Failed to execute command: {str(e)}")
|
||||
|
||||
def get_router_info(self) -> Dict[str, str]:
|
||||
"""Get router system information.
|
||||
|
||||
Returns:
|
||||
Dictionary containing router information
|
||||
"""
|
||||
# Try common commands to get router info
|
||||
info = {}
|
||||
|
||||
try:
|
||||
# Try to get model information
|
||||
model_output = self.execute_command("cat /proc/cpuinfo | grep 'model name' | head -1")
|
||||
if model_output:
|
||||
info['model'] = model_output.split(':')[-1].strip()
|
||||
else:
|
||||
info['model'] = "Unknown"
|
||||
except:
|
||||
info['model'] = "Unknown"
|
||||
|
||||
try:
|
||||
# Try to get firmware version
|
||||
firmware_output = self.execute_command("cat /etc/openwrt_release | grep DISTRIB_RELEASE")
|
||||
if firmware_output:
|
||||
info['firmware'] = firmware_output.split('=')[-1].strip().strip("'\"")
|
||||
else:
|
||||
info['firmware'] = "Unknown"
|
||||
except:
|
||||
info['firmware'] = "Unknown"
|
||||
|
||||
return info
|
||||
|
||||
def enable_monitor_mode(self, interface: str) -> bool:
|
||||
"""Enable monitor mode on WiFi interface.
|
||||
|
||||
Args:
|
||||
interface: WiFi interface name (e.g., 'wlan0')
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Bring interface down
|
||||
self.execute_command(f"ifconfig {interface} down")
|
||||
|
||||
# Set monitor mode
|
||||
self.execute_command(f"iwconfig {interface} mode monitor")
|
||||
|
||||
# Bring interface up
|
||||
self.execute_command(f"ifconfig {interface} up")
|
||||
|
||||
return True
|
||||
|
||||
except RouterConnectionError:
|
||||
return False
|
||||
|
||||
def disable_monitor_mode(self, interface: str) -> bool:
|
||||
"""Disable monitor mode on WiFi interface.
|
||||
|
||||
Args:
|
||||
interface: WiFi interface name (e.g., 'wlan0')
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Bring interface down
|
||||
self.execute_command(f"ifconfig {interface} down")
|
||||
|
||||
# Set managed mode
|
||||
self.execute_command(f"iwconfig {interface} mode managed")
|
||||
|
||||
# Bring interface up
|
||||
self.execute_command(f"ifconfig {interface} up")
|
||||
|
||||
return True
|
||||
|
||||
except RouterConnectionError:
|
||||
return False
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit."""
|
||||
self.disconnect()
|
||||
@@ -0,0 +1,279 @@
|
||||
"""DensePose head for WiFi-DensePose system."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict, Any, Tuple, List
|
||||
|
||||
|
||||
class DensePoseError(Exception):
|
||||
"""Exception raised for DensePose head errors."""
|
||||
pass
|
||||
|
||||
|
||||
class DensePoseHead(nn.Module):
|
||||
"""DensePose head for body part segmentation and UV coordinate regression."""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""Initialize DensePose head.
|
||||
|
||||
Args:
|
||||
config: Configuration dictionary with head parameters
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self._validate_config(config)
|
||||
self.config = config
|
||||
|
||||
self.input_channels = config['input_channels']
|
||||
self.num_body_parts = config['num_body_parts']
|
||||
self.num_uv_coordinates = config['num_uv_coordinates']
|
||||
self.hidden_channels = config.get('hidden_channels', [128, 64])
|
||||
self.kernel_size = config.get('kernel_size', 3)
|
||||
self.padding = config.get('padding', 1)
|
||||
self.dropout_rate = config.get('dropout_rate', 0.1)
|
||||
self.use_deformable_conv = config.get('use_deformable_conv', False)
|
||||
self.use_fpn = config.get('use_fpn', False)
|
||||
self.fpn_levels = config.get('fpn_levels', [2, 3, 4, 5])
|
||||
self.output_stride = config.get('output_stride', 4)
|
||||
|
||||
# Feature Pyramid Network (optional)
|
||||
if self.use_fpn:
|
||||
self.fpn = self._build_fpn()
|
||||
|
||||
# Shared feature processing
|
||||
self.shared_conv = self._build_shared_layers()
|
||||
|
||||
# Segmentation head for body part classification
|
||||
self.segmentation_head = self._build_segmentation_head()
|
||||
|
||||
# UV regression head for coordinate prediction
|
||||
self.uv_regression_head = self._build_uv_regression_head()
|
||||
|
||||
# Initialize weights
|
||||
self._initialize_weights()
|
||||
|
||||
def _validate_config(self, config: Dict[str, Any]):
|
||||
"""Validate configuration parameters."""
|
||||
required_fields = ['input_channels', 'num_body_parts', 'num_uv_coordinates']
|
||||
for field in required_fields:
|
||||
if field not in config:
|
||||
raise ValueError(f"Missing required field: {field}")
|
||||
|
||||
if config['input_channels'] <= 0:
|
||||
raise ValueError("input_channels must be positive")
|
||||
|
||||
if config['num_body_parts'] <= 0:
|
||||
raise ValueError("num_body_parts must be positive")
|
||||
|
||||
if config['num_uv_coordinates'] <= 0:
|
||||
raise ValueError("num_uv_coordinates must be positive")
|
||||
|
||||
def _build_fpn(self) -> nn.Module:
|
||||
"""Build Feature Pyramid Network."""
|
||||
return nn.ModuleDict({
|
||||
f'level_{level}': nn.Conv2d(self.input_channels, self.input_channels, 1)
|
||||
for level in self.fpn_levels
|
||||
})
|
||||
|
||||
def _build_shared_layers(self) -> nn.Module:
|
||||
"""Build shared feature processing layers."""
|
||||
layers = []
|
||||
in_channels = self.input_channels
|
||||
|
||||
for hidden_dim in self.hidden_channels:
|
||||
layers.extend([
|
||||
nn.Conv2d(in_channels, hidden_dim,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.padding),
|
||||
nn.BatchNorm2d(hidden_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout2d(self.dropout_rate)
|
||||
])
|
||||
in_channels = hidden_dim
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _build_segmentation_head(self) -> nn.Module:
|
||||
"""Build segmentation head for body part classification."""
|
||||
final_hidden = self.hidden_channels[-1] if self.hidden_channels else self.input_channels
|
||||
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(final_hidden, final_hidden // 2,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.padding),
|
||||
nn.BatchNorm2d(final_hidden // 2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout2d(self.dropout_rate),
|
||||
|
||||
# Upsampling to increase resolution
|
||||
nn.ConvTranspose2d(final_hidden // 2, final_hidden // 4,
|
||||
kernel_size=4, stride=2, padding=1),
|
||||
nn.BatchNorm2d(final_hidden // 4),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(final_hidden // 4, self.num_body_parts + 1, kernel_size=1),
|
||||
# +1 for background class
|
||||
)
|
||||
|
||||
def _build_uv_regression_head(self) -> nn.Module:
|
||||
"""Build UV regression head for coordinate prediction."""
|
||||
final_hidden = self.hidden_channels[-1] if self.hidden_channels else self.input_channels
|
||||
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(final_hidden, final_hidden // 2,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.padding),
|
||||
nn.BatchNorm2d(final_hidden // 2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout2d(self.dropout_rate),
|
||||
|
||||
# Upsampling to increase resolution
|
||||
nn.ConvTranspose2d(final_hidden // 2, final_hidden // 4,
|
||||
kernel_size=4, stride=2, padding=1),
|
||||
nn.BatchNorm2d(final_hidden // 4),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(final_hidden // 4, self.num_uv_coordinates, kernel_size=1),
|
||||
)
|
||||
|
||||
def _initialize_weights(self):
|
||||
"""Initialize network weights."""
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
|
||||
"""Forward pass through the DensePose head.
|
||||
|
||||
Args:
|
||||
x: Input feature tensor of shape (batch_size, channels, height, width)
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- segmentation: Body part logits (batch_size, num_parts+1, height, width)
|
||||
- uv_coordinates: UV coordinates (batch_size, 2, height, width)
|
||||
"""
|
||||
# Validate input shape
|
||||
if x.shape[1] != self.input_channels:
|
||||
raise DensePoseError(f"Expected {self.input_channels} input channels, got {x.shape[1]}")
|
||||
|
||||
# Apply FPN if enabled
|
||||
if self.use_fpn:
|
||||
# Simple FPN processing - in practice this would be more sophisticated
|
||||
x = self.fpn['level_2'](x)
|
||||
|
||||
# Shared feature processing
|
||||
shared_features = self.shared_conv(x)
|
||||
|
||||
# Segmentation branch
|
||||
segmentation_logits = self.segmentation_head(shared_features)
|
||||
|
||||
# UV regression branch
|
||||
uv_coordinates = self.uv_regression_head(shared_features)
|
||||
uv_coordinates = torch.sigmoid(uv_coordinates) # Normalize to [0, 1]
|
||||
|
||||
return {
|
||||
'segmentation': segmentation_logits,
|
||||
'uv_coordinates': uv_coordinates
|
||||
}
|
||||
|
||||
def compute_segmentation_loss(self, pred_logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute segmentation loss.
|
||||
|
||||
Args:
|
||||
pred_logits: Predicted segmentation logits
|
||||
target: Target segmentation masks
|
||||
|
||||
Returns:
|
||||
Computed cross-entropy loss
|
||||
"""
|
||||
return F.cross_entropy(pred_logits, target, ignore_index=-1)
|
||||
|
||||
def compute_uv_loss(self, pred_uv: torch.Tensor, target_uv: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute UV coordinate regression loss.
|
||||
|
||||
Args:
|
||||
pred_uv: Predicted UV coordinates
|
||||
target_uv: Target UV coordinates
|
||||
|
||||
Returns:
|
||||
Computed L1 loss
|
||||
"""
|
||||
return F.l1_loss(pred_uv, target_uv)
|
||||
|
||||
def compute_total_loss(self, predictions: Dict[str, torch.Tensor],
|
||||
seg_target: torch.Tensor,
|
||||
uv_target: torch.Tensor,
|
||||
seg_weight: float = 1.0,
|
||||
uv_weight: float = 1.0) -> torch.Tensor:
|
||||
"""Compute total loss combining segmentation and UV losses.
|
||||
|
||||
Args:
|
||||
predictions: Dictionary of predictions
|
||||
seg_target: Target segmentation masks
|
||||
uv_target: Target UV coordinates
|
||||
seg_weight: Weight for segmentation loss
|
||||
uv_weight: Weight for UV loss
|
||||
|
||||
Returns:
|
||||
Combined loss
|
||||
"""
|
||||
seg_loss = self.compute_segmentation_loss(predictions['segmentation'], seg_target)
|
||||
uv_loss = self.compute_uv_loss(predictions['uv_coordinates'], uv_target)
|
||||
|
||||
return seg_weight * seg_loss + uv_weight * uv_loss
|
||||
|
||||
def get_prediction_confidence(self, predictions: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||||
"""Get prediction confidence scores.
|
||||
|
||||
Args:
|
||||
predictions: Dictionary of predictions
|
||||
|
||||
Returns:
|
||||
Dictionary of confidence scores
|
||||
"""
|
||||
seg_logits = predictions['segmentation']
|
||||
uv_coords = predictions['uv_coordinates']
|
||||
|
||||
# Segmentation confidence: max probability
|
||||
seg_probs = F.softmax(seg_logits, dim=1)
|
||||
seg_confidence = torch.max(seg_probs, dim=1)[0]
|
||||
|
||||
# UV confidence: inverse of prediction variance
|
||||
uv_variance = torch.var(uv_coords, dim=1, keepdim=True)
|
||||
uv_confidence = 1.0 / (1.0 + uv_variance)
|
||||
|
||||
return {
|
||||
'segmentation_confidence': seg_confidence,
|
||||
'uv_confidence': uv_confidence.squeeze(1)
|
||||
}
|
||||
|
||||
def post_process_predictions(self, predictions: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||||
"""Post-process predictions for final output.
|
||||
|
||||
Args:
|
||||
predictions: Raw predictions from forward pass
|
||||
|
||||
Returns:
|
||||
Post-processed predictions
|
||||
"""
|
||||
seg_logits = predictions['segmentation']
|
||||
uv_coords = predictions['uv_coordinates']
|
||||
|
||||
# Convert logits to class predictions
|
||||
body_parts = torch.argmax(seg_logits, dim=1)
|
||||
|
||||
# Get confidence scores
|
||||
confidence = self.get_prediction_confidence(predictions)
|
||||
|
||||
return {
|
||||
'body_parts': body_parts,
|
||||
'uv_coordinates': uv_coords,
|
||||
'confidence_scores': confidence
|
||||
}
|
||||
@@ -3,7 +3,12 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict, Any
|
||||
from typing import Dict, Any, List
|
||||
|
||||
|
||||
class ModalityTranslationError(Exception):
|
||||
"""Exception raised for modality translation errors."""
|
||||
pass
|
||||
|
||||
|
||||
class ModalityTranslationNetwork(nn.Module):
|
||||
@@ -17,11 +22,20 @@ class ModalityTranslationNetwork(nn.Module):
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self._validate_config(config)
|
||||
self.config = config
|
||||
|
||||
self.input_channels = config['input_channels']
|
||||
self.hidden_dim = config['hidden_dim']
|
||||
self.output_dim = config['output_dim']
|
||||
self.num_layers = config['num_layers']
|
||||
self.dropout_rate = config['dropout_rate']
|
||||
self.hidden_channels = config['hidden_channels']
|
||||
self.output_channels = config['output_channels']
|
||||
self.kernel_size = config.get('kernel_size', 3)
|
||||
self.stride = config.get('stride', 1)
|
||||
self.padding = config.get('padding', 1)
|
||||
self.dropout_rate = config.get('dropout_rate', 0.1)
|
||||
self.activation = config.get('activation', 'relu')
|
||||
self.normalization = config.get('normalization', 'batch')
|
||||
self.use_attention = config.get('use_attention', False)
|
||||
self.attention_heads = config.get('attention_heads', 8)
|
||||
|
||||
# Encoder: CSI -> Feature space
|
||||
self.encoder = self._build_encoder()
|
||||
@@ -29,57 +43,114 @@ class ModalityTranslationNetwork(nn.Module):
|
||||
# Decoder: Feature space -> Visual-like features
|
||||
self.decoder = self._build_decoder()
|
||||
|
||||
# Attention mechanism
|
||||
if self.use_attention:
|
||||
self.attention = self._build_attention()
|
||||
|
||||
# Initialize weights
|
||||
self._initialize_weights()
|
||||
|
||||
def _build_encoder(self) -> nn.Module:
|
||||
def _validate_config(self, config: Dict[str, Any]):
|
||||
"""Validate configuration parameters."""
|
||||
required_fields = ['input_channels', 'hidden_channels', 'output_channels']
|
||||
for field in required_fields:
|
||||
if field not in config:
|
||||
raise ValueError(f"Missing required field: {field}")
|
||||
|
||||
if config['input_channels'] <= 0:
|
||||
raise ValueError("input_channels must be positive")
|
||||
|
||||
if not config['hidden_channels'] or len(config['hidden_channels']) == 0:
|
||||
raise ValueError("hidden_channels must be a non-empty list")
|
||||
|
||||
if config['output_channels'] <= 0:
|
||||
raise ValueError("output_channels must be positive")
|
||||
|
||||
def _build_encoder(self) -> nn.ModuleList:
|
||||
"""Build encoder network."""
|
||||
layers = []
|
||||
layers = nn.ModuleList()
|
||||
|
||||
# Initial convolution
|
||||
layers.append(nn.Conv2d(self.input_channels, 64, kernel_size=3, padding=1))
|
||||
layers.append(nn.BatchNorm2d(64))
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
layers.append(nn.Dropout2d(self.dropout_rate))
|
||||
in_channels = self.input_channels
|
||||
|
||||
# Progressive downsampling
|
||||
in_channels = 64
|
||||
for i in range(self.num_layers - 1):
|
||||
out_channels = min(in_channels * 2, self.hidden_dim)
|
||||
layers.extend([
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
for i, out_channels in enumerate(self.hidden_channels):
|
||||
layer_block = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride if i == 0 else 2,
|
||||
padding=self.padding),
|
||||
self._get_normalization(out_channels),
|
||||
self._get_activation(),
|
||||
nn.Dropout2d(self.dropout_rate)
|
||||
])
|
||||
)
|
||||
layers.append(layer_block)
|
||||
in_channels = out_channels
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
return layers
|
||||
|
||||
def _build_decoder(self) -> nn.Module:
|
||||
def _build_decoder(self) -> nn.ModuleList:
|
||||
"""Build decoder network."""
|
||||
layers = []
|
||||
layers = nn.ModuleList()
|
||||
|
||||
# Get the actual output channels from encoder (should be hidden_dim)
|
||||
encoder_out_channels = self.hidden_dim
|
||||
# Start with the last hidden channel size
|
||||
in_channels = self.hidden_channels[-1]
|
||||
|
||||
# Progressive upsampling
|
||||
in_channels = encoder_out_channels
|
||||
for i in range(self.num_layers - 1):
|
||||
out_channels = max(in_channels // 2, 64)
|
||||
layers.extend([
|
||||
nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
# Progressive upsampling (reverse of encoder)
|
||||
for i, out_channels in enumerate(reversed(self.hidden_channels[:-1])):
|
||||
layer_block = nn.Sequential(
|
||||
nn.ConvTranspose2d(in_channels, out_channels,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=2,
|
||||
padding=self.padding,
|
||||
output_padding=1),
|
||||
self._get_normalization(out_channels),
|
||||
self._get_activation(),
|
||||
nn.Dropout2d(self.dropout_rate)
|
||||
])
|
||||
)
|
||||
layers.append(layer_block)
|
||||
in_channels = out_channels
|
||||
|
||||
# Final output layer
|
||||
layers.append(nn.Conv2d(in_channels, self.output_dim, kernel_size=3, padding=1))
|
||||
layers.append(nn.Tanh()) # Normalize output
|
||||
final_layer = nn.Sequential(
|
||||
nn.Conv2d(in_channels, self.output_channels,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.padding),
|
||||
nn.Tanh() # Normalize output
|
||||
)
|
||||
layers.append(final_layer)
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
return layers
|
||||
|
||||
def _get_normalization(self, channels: int) -> nn.Module:
|
||||
"""Get normalization layer."""
|
||||
if self.normalization == 'batch':
|
||||
return nn.BatchNorm2d(channels)
|
||||
elif self.normalization == 'instance':
|
||||
return nn.InstanceNorm2d(channels)
|
||||
elif self.normalization == 'layer':
|
||||
return nn.GroupNorm(1, channels)
|
||||
else:
|
||||
return nn.Identity()
|
||||
|
||||
def _get_activation(self) -> nn.Module:
|
||||
"""Get activation function."""
|
||||
if self.activation == 'relu':
|
||||
return nn.ReLU(inplace=True)
|
||||
elif self.activation == 'leaky_relu':
|
||||
return nn.LeakyReLU(0.2, inplace=True)
|
||||
elif self.activation == 'gelu':
|
||||
return nn.GELU()
|
||||
else:
|
||||
return nn.ReLU(inplace=True)
|
||||
|
||||
def _build_attention(self) -> nn.Module:
|
||||
"""Build attention mechanism."""
|
||||
return nn.MultiheadAttention(
|
||||
embed_dim=self.hidden_channels[-1],
|
||||
num_heads=self.attention_heads,
|
||||
dropout=self.dropout_rate,
|
||||
batch_first=True
|
||||
)
|
||||
|
||||
def _initialize_weights(self):
|
||||
"""Initialize network weights."""
|
||||
@@ -103,12 +174,128 @@ class ModalityTranslationNetwork(nn.Module):
|
||||
"""
|
||||
# Validate input shape
|
||||
if x.shape[1] != self.input_channels:
|
||||
raise RuntimeError(f"Expected {self.input_channels} input channels, got {x.shape[1]}")
|
||||
raise ModalityTranslationError(f"Expected {self.input_channels} input channels, got {x.shape[1]}")
|
||||
|
||||
# Encode CSI data
|
||||
encoded = self.encoder(x)
|
||||
encoded_features = self.encode(x)
|
||||
|
||||
# Decode to visual-like features
|
||||
decoded = self.decoder(encoded)
|
||||
decoded = self.decode(encoded_features)
|
||||
|
||||
return decoded
|
||||
return decoded
|
||||
|
||||
def encode(self, x: torch.Tensor) -> List[torch.Tensor]:
|
||||
"""Encode input through encoder layers.
|
||||
|
||||
Args:
|
||||
x: Input tensor
|
||||
|
||||
Returns:
|
||||
List of feature maps from each encoder layer
|
||||
"""
|
||||
features = []
|
||||
current = x
|
||||
|
||||
for layer in self.encoder:
|
||||
current = layer(current)
|
||||
features.append(current)
|
||||
|
||||
return features
|
||||
|
||||
def decode(self, encoded_features: List[torch.Tensor]) -> torch.Tensor:
|
||||
"""Decode features through decoder layers.
|
||||
|
||||
Args:
|
||||
encoded_features: List of encoded feature maps
|
||||
|
||||
Returns:
|
||||
Decoded output tensor
|
||||
"""
|
||||
# Start with the last encoded feature
|
||||
current = encoded_features[-1]
|
||||
|
||||
# Apply attention if enabled
|
||||
if self.use_attention:
|
||||
batch_size, channels, height, width = current.shape
|
||||
# Reshape for attention: (batch, seq_len, embed_dim)
|
||||
current_flat = current.view(batch_size, channels, -1).transpose(1, 2)
|
||||
attended, _ = self.attention(current_flat, current_flat, current_flat)
|
||||
current = attended.transpose(1, 2).view(batch_size, channels, height, width)
|
||||
|
||||
# Apply decoder layers
|
||||
for layer in self.decoder:
|
||||
current = layer(current)
|
||||
|
||||
return current
|
||||
|
||||
def compute_translation_loss(self, predicted: torch.Tensor, target: torch.Tensor, loss_type: str = 'mse') -> torch.Tensor:
|
||||
"""Compute translation loss between predicted and target features.
|
||||
|
||||
Args:
|
||||
predicted: Predicted feature tensor
|
||||
target: Target feature tensor
|
||||
loss_type: Type of loss ('mse', 'l1', 'smooth_l1')
|
||||
|
||||
Returns:
|
||||
Computed loss tensor
|
||||
"""
|
||||
if loss_type == 'mse':
|
||||
return F.mse_loss(predicted, target)
|
||||
elif loss_type == 'l1':
|
||||
return F.l1_loss(predicted, target)
|
||||
elif loss_type == 'smooth_l1':
|
||||
return F.smooth_l1_loss(predicted, target)
|
||||
else:
|
||||
return F.mse_loss(predicted, target)
|
||||
|
||||
def get_feature_statistics(self, features: torch.Tensor) -> Dict[str, float]:
|
||||
"""Get statistics of feature tensor.
|
||||
|
||||
Args:
|
||||
features: Feature tensor to analyze
|
||||
|
||||
Returns:
|
||||
Dictionary of feature statistics
|
||||
"""
|
||||
with torch.no_grad():
|
||||
return {
|
||||
'mean': features.mean().item(),
|
||||
'std': features.std().item(),
|
||||
'min': features.min().item(),
|
||||
'max': features.max().item(),
|
||||
'sparsity': (features == 0).float().mean().item()
|
||||
}
|
||||
|
||||
def get_intermediate_features(self, x: torch.Tensor) -> Dict[str, Any]:
|
||||
"""Get intermediate features for visualization.
|
||||
|
||||
Args:
|
||||
x: Input tensor
|
||||
|
||||
Returns:
|
||||
Dictionary containing intermediate features
|
||||
"""
|
||||
result = {}
|
||||
|
||||
# Get encoder features
|
||||
encoder_features = self.encode(x)
|
||||
result['encoder_features'] = encoder_features
|
||||
|
||||
# Get decoder features
|
||||
decoder_features = []
|
||||
current = encoder_features[-1]
|
||||
|
||||
if self.use_attention:
|
||||
batch_size, channels, height, width = current.shape
|
||||
current_flat = current.view(batch_size, channels, -1).transpose(1, 2)
|
||||
attended, attention_weights = self.attention(current_flat, current_flat, current_flat)
|
||||
current = attended.transpose(1, 2).view(batch_size, channels, height, width)
|
||||
result['attention_weights'] = attention_weights
|
||||
|
||||
for layer in self.decoder:
|
||||
current = layer(current)
|
||||
decoder_features.append(current)
|
||||
|
||||
result['decoder_features'] = decoder_features
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user