xttmp 2.3.0__py3-none-any.whl

xttmp/core/dstmd_core.py

@@ -0,0 +1,213 @@
+from collections import deque
+
+import torch
+import torch.nn.functional as F
+
+from .base_core import BaseCore
+from .estmd_core import Tm1, Mi1
+from .math_operator import GammaDelay, SpatialInhibition
+from ..util.create_kernel import create_direction_inhi_kernel
+
+
+class Medulla(BaseCore):
+    """Medulla class for motion detection."""
+    
+    def __init__(self):
+        """Constructor method."""
+        # Initializes the Medulla object
+        super().__init__()
+
+        # Initialize components
+        self.tm3 = Tm3()
+        self.mi1_para4 = Mi1(3, 15)
+
+        self.tm2 = Tm2()
+        self.tm1_para5 = Tm1(5, 25)
+        self.tm1_para6 = Tm1(8, 40)
+        
+    def setup(self):
+        """Initialization method."""
+        # Initializes the delay components
+
+        self.mi1_para4.setup()
+        self.tm1_para5.setup()
+        self.tm1_para6.setup()
+
+    def forward(self, x):
+        """Processing method."""
+        # Processes input signals and produces output
+        
+        # Process Tm3 and Tm2 signals
+        tm3_output = self.tm3.forward(x) # L ON
+        tm2_output = self.tm2.forward(x) # L OFF
+        
+        # Process signals with delays
+        mi1_para4_output = self.mi1_para4.forward(tm3_output)
+        tm1_para5_output = self.tm1_para5.forward(tm2_output)
+        tm1_para6_output = self.tm1_para6.forward(tm2_output)
+
+        # Output signals
+        self.output = (tm3_output, mi1_para4_output, tm1_para5_output, tm1_para6_output)
+        return self.output
+
+
+class Tm2(BaseCore):
+    """Tm2 class for motion detection."""
+
+    def forward(self, iptMatrix):
+        """Processing method."""
+        # Processes the input matrix by performing a maximum operation with zero for negative values
+        self.output = torch.clamp(-iptMatrix, min=0)
+
+        return self.output
+
+
+class Tm3(BaseCore):
+    """Tm3 class for motion detection."""
+
+    def forward(self, iptMatrix):
+        """Processing method."""
+        # Processes the input matrix by performing a maximum operation with zero for negative values
+
+        self.output = torch.clamp(iptMatrix, min=0)
+
+        return self.output
+
+
+class Lobula(BaseCore):
+    """Lobula class for motion detection."""
+
+    def __init__(self):
+        """Constructor method."""
+        # Initializes the Lobula object
+        super().__init__()
+        self.alpha1 = 3  # Alpha parameter
+
+        self.register_buffer('theta_list', torch.tensor([(i * torch.pi / 4) for i in range(8)]))
+        self.hLateralInhi = SpatialInhibition()  # Lateral inhibition component
+        self.hDirectionInhi = DirectionInhibition()  # Directional inhibition component
+
+    def setup(self):
+        """Initialization method."""
+        # Initializes the lateral and directional inhibition components
+        self.hLateralInhi.setup()
+        self.hDirectionInhi.setup()
+
+    def forward(self, tm3, mi1_p4, tm1_p5, tm1_p6):
+        # tm3, mi1_p4, tm1_p5, tm1_p6 形状均为 [1, 1, H, W]
+        _, _, imgH, imgW = tm3.shape
+        device = tm3.device
+
+        num_thetas = len(self.theta_list)
+        a1 = self.alpha1
+
+        # 2. 计算偏移索引
+        shifts_x = torch.round(a1 * torch.cos(self.theta_list)).long() # 形状: [num_thetas]
+        shifts_y = torch.round(a1 * torch.sin(self.theta_list)).long() # 形状: [num_thetas]
+
+        # 3. 提取中心 ROI
+        y_s, y_e = a1, imgH - a1
+        x_s, x_e = a1, imgW - a1
+        
+        # 提取不变部分的 ROI，保持 4D 形状: [1, 1, h_roi, w_roi]
+        tm3_roi = tm3[:, :, y_s:y_e, x_s:x_e]
+        tm1_p5_roi = tm1_p5[:, :, y_s:y_e, x_s:x_e]
+
+        # 4. 生成偏移索引网格 (关键点)
+        grid_y, grid_x = torch.meshgrid(
+            torch.arange(y_s, y_e, device=device),
+            torch.arange(x_s, x_e, device=device),
+            indexing='ij'
+        )
+
+        # 计算所有方向的索引: [num_thetas, h_roi, w_roi]
+        src_idx_x = grid_x.unsqueeze(0).unsqueeze(0) - shifts_x .view(1, -1, 1, 1)
+        src_idx_y = grid_y.unsqueeze(0).unsqueeze(0) + shifts_y.view(1, -1, 1, 1)
+
+        # 5. 高级索引提取偏移信号
+        # mi1_p4[0, 0] 是 [H, W]，通过 src_idx 提取后变成 [num_thetas, h_roi, w_roi]
+        # 我们将其扩展回 4D: [1, num_thetas, h_roi, w_roi]
+        mi1_p4_shifted = mi1_p4[0, 0, src_idx_y, src_idx_x]
+        tm1_p6_shifted = tm1_p6[0, 0, src_idx_y, src_idx_x]
+
+        # 6. 计算相关输出 (利用广播)
+        # tm3_roi: [1, 1, h_roi, w_roi]
+        # mi1_p4_shifted: [1, num_thetas, h_roi, w_roi]
+        # 结果 corre_roi: [1, num_thetas, h_roi, w_roi]
+        corre_roi = tm3_roi * (tm1_p5_roi + mi1_p4_shifted) * tm1_p6_shifted
+
+        # 7. 填回全零张量
+        correOutput = torch.zeros((1, num_thetas, imgH, imgW), device=device)
+        correOutput[:, :, y_s:y_e, x_s:x_e] = corre_roi
+
+        lateralInhioutput = self.hLateralInhi.forward(correOutput)
+        self.output = self.hDirectionInhi.forward(lateralInhioutput)
+
+        return self.output
+
+
+class DirectionInhibition(BaseCore):
+    """Directional inhibition in DSTMD."""
+
+    def __init__(self):
+        """Constructor method."""
+        # Initializes the DirectionInhi object
+        super().__init__()
+        self.direction = 8  # Number of directions
+        self.sigma1 = 1.5  # Sigma for the first Gaussian kernel
+        self.sigma2 = 3.0  # Sigma for the second Gaussian kernel
+
+        self.register_buffer("diretional_inhi_kernel", torch.empty(0))  # Placeholder for the directional inhibition kernel
+
+        self.setup()  # Initialize the kernel
+
+    def setup(self):
+        """Initialization method."""
+        # Initializes the directional inhibition kernel
+
+        _diretional_inhi_kernel = create_direction_inhi_kernel(
+                self.direction, self.sigma1, self.sigma2
+            )
+        # Shape: [1, 1, kernel_size]
+        self.diretional_inhi_kernel.data = _diretional_inhi_kernel
+
+    def forward(self, x):
+        """
+            Input x shape: [B, C, H, W], where C = self.direction
+            Output shape: [B, C, H, W]        
+        """
+        # Performs directional inhibition on the input
+
+        # 1. 准备数据维度
+        # 输入是 [B, C, H, W]，卷积需要在 C 维度上滑，所以要把 H, W 暂时视为 Batch
+        b, c, h, w = x.shape
+        
+        # 转换形状: [B, C, H, W] -> [B, H, W, C] -> [B*H*W, C]
+        # 这样对于 conv1d 来说，BatchSize = B*H*W, 通道数 = 1, 序列长度 = C
+        x = x.permute(0, 2, 3, 1).reshape(b * h * w, 1, c)
+        
+        # 3. 执行循环卷积 (Circular Convolution)
+        # padding 设为 kernel_size // 2，且模式设为 'circular'
+        pad_size = self.diretional_inhi_kernel.shape[-1]
+        center_idx = pad_size // 2
+        
+        pad_left = center_idx
+        pad_right = pad_size - center_idx - 1
+        
+        # F.pad 在 1D 信号上的填充格式是 (left, right)
+        x_padded = F.pad(x, (pad_left, pad_right), mode='circular')
+
+        # F.conv1d 会在 C 维度（方向轴）上滑动
+        # 结果形状依然是 [B*H*W, 1, C]
+        result = F.conv1d(x_padded, self.diretional_inhi_kernel)
+        
+        # 4. 激活与恢复形状
+        opt = F.relu(result)
+        
+        # 恢复回 [B, C, H, W]
+        return opt.reshape(b, h, w, c).permute(0, 3, 1, 2)
+
+
+    
+
+

xttmp/core/estmd_backbone.py

@@ -0,0 +1,110 @@
+import torch
+
+from .base_core import BaseCore
+from .math_operator import GammaBandPassFilter, SpatialInhibition
+from . import estmd_core
+
+
+class Lamina(GammaBandPassFilter):
+    pass
+
+
+class Medulla(BaseCore):
+    """Medulla layer of the motion detection system."""
+    
+    def __init__(self):
+        """Constructor method."""
+        # Initializes the Medulla object
+        super().__init__()
+        # Initialize components
+        self.tm1 = estmd_core.Tm1()
+        self.mi1 = estmd_core.Mi1()
+        self.tm2 = Tm2()
+        self.tm3 = Tm3()
+
+    def setup(self):
+        """Initialization method."""
+        # This method initializes the Medulla layer components
+        self.tm1.setup()
+        self.mi1.setup()
+        self.tm2.setup()
+        self.tm3.setup()
+
+    def forward(self, MedullaIpt):
+        """Processing method."""
+        # Applies processing to the input and returns the output
+        
+        # Process Tm2 and Tm3 components
+        tm2_output = self.tm2.forward(MedullaIpt)
+        tm3_output = self.tm3.forward(MedullaIpt)
+
+        # Process Tm1 component using output of Tm2
+        tm1_output = self.tm1.forward(tm2_output)
+
+        # Store the output signals in output property
+        self.output = [tm3_output, tm1_output]
+
+        return self.output
+
+
+class Lobula(BaseCore):
+    """Lobula layer of the motion detection system."""
+    
+    def __init__(self):
+        """Constructor method."""
+        # Initializes the Lobula object
+        super().__init__()
+        # Initialize the SpatialInhibition component
+        self.spatial_inhibition = SpatialInhibition()
+        # Parameters related to the recombination of ON and OFF channels
+        self.a = 0
+        self.b = 0
+        self.c = 1
+
+    def setup(self):
+        """Initialization method."""
+        # This method initializes the Lobula layer component
+        self.spatial_inhibition.setup()
+
+    def forward(self, on_output, off_output):
+        """Processing method."""
+        # Performs a correlation operation on the ON and OFF channels
+        # and then applies surround inhibition
+        
+        # Perform the correlation operation
+        correlationOutput = (
+            self.a * on_output +
+            self.b * off_output +
+            self.c * on_output * off_output
+        )
+
+        # Apply surround inhibition
+        lobulaoutput = self.spatial_inhibition.forward(correlationOutput)
+        
+        # Store the output in output property
+        self.output = lobulaoutput
+        return lobulaoutput, correlationOutput
+
+
+class Tm2(BaseCore):
+    """Tm2 without spatial inhibition."""
+    def forward(self, x):
+        """Processing method."""
+        # Applies surround inhibition to the input to generate the output
+        self.output = torch.clamp(-x, min=0)  # Apply surround inhibition
+        return self.output
+    
+
+class Tm3(BaseCore):
+    """Tm3 without spatial inhibition."""
+    def forward(self, tm3outputIpt):
+        """Processing method."""
+        # Applies a surround inhibition to the input to generate the output
+        self.output = torch.clamp(tm3outputIpt, min=0)  # Apply surround inhibition
+        return self.output
+
+
+
+
+
+

xttmp/core/estmd_core.py

@@ -0,0 +1,356 @@
+from collections import deque
+
+import torch
+import torch.nn.functional as F
+
+from .base_core import BaseCore
+from .math_operator import (compute_temporal_conv_inplace,
+                            GaussianBlur, SpatialInhibition, 
+                            GammaDelay, GammaBandPassFilter)
+from ..util.create_kernel import create_2d_gaussian_kernel
+
+
+class Retina(GaussianBlur):
+    pass
+
+
+class Lamina(BaseCore):
+    """
+    LAMINA Lamina layer
+    This class implements the Lamina layer of the ESTMD
+    
+    Author: Mingshuo Xu
+    Date: 2024-04-29
+    """
+    
+    def __init__(self):
+        """
+        Constructor
+        Initializes the Lamina object and creates GammaBankPassFilter
+        and LaminaLateralInhibition objects
+        """
+        super().__init__()
+        self.gamma_BPF = GammaBandPassFilter()
+        self.spatial_inhibition = LaminaLateralInhibition()
+
+    def setup(self):
+        self.gamma_BPF.setup()
+        self.spatial_inhibition.setup()
+
+    def reset_buffer(self):
+        self.gamma_BPF.reset_buffer()
+
+    def forward(self, laminaIpt):
+        """
+        Processing method
+        Applies GammaBankPassFilter and LaminaLateralInhibition to the input matrix
+        
+        Parameters:
+        - laminaIpt: Input matrix
+        
+        Returns:
+        - laminaOpt: Processed output matrix
+        """
+        signalWithBPF = self.gamma_BPF.forward(laminaIpt)
+        self.output = self.spatial_inhibition.forward(signalWithBPF)
+
+        return self.output
+
+
+class Medulla(BaseCore):
+    """
+    Medulla Layer of DSTMD
+    This class implements the Medulla layer of the ESTMD.
+    """
+    
+    def __init__(self):
+        """
+        Constructor method
+        Initializes the Medulla object
+        """
+        super().__init__()
+        self.tm1 = Tm1(order=12, tau=25)  # Initialize Tm1 object
+        self.tm2 = Tm2()  # Initialize Tm2 object
+        self.tm3 = Tm3()  # Initialize Tm3 object
+        self.mi1 = Mi1(order=12, tau=25)  # Initialize Tm3 object
+
+    def setup(self):
+        """
+        Initialization method
+        Initializes the Tm1, Tm2, and Tm3 objects
+        """
+        self.tm1.setup()
+        self.tm2.setup()
+        self.tm3.setup()
+        self.mi1.setup()
+
+    def reset_buffer(self):
+        """
+        Buffer reset method
+        Resets the buffers of Tm1, Tm2, and Tm3 objects
+        """
+        self.tm1.reset_buffer()
+        self.mi1.reset_buffer()
+
+    def forward(self, x):
+        """
+        Processing method
+        Processes the input MedullaIpt through Tm1, Tm2, and Tm3 layers
+        
+        Parameters:
+        - MedullaIpt: Input matrix
+        
+        Returns:
+        - tm3Signal: Output of Tm3 layer
+        - tm1Signal: Output of Tm1 layer
+        """
+        tm2_output = self.tm2.forward(x)  # Process input through Tm2
+        tm3_output = self.tm3.forward(x)  # Process input through Tm3
+
+        tm1_output = self.tm1.forward(tm2_output)  # Process Tm2 output through Tm1
+
+        self.output = (tm3_output, tm1_output)  # Update output property with output
+        return tm3_output, tm1_output
+
+
+class Lobula(BaseCore):
+    """
+    Lobula Layer of DSTMD
+    This class implements the Lobula layer of the ESTMD.
+    """
+    
+    def __init__(self):
+        """
+        Constructor method
+        Initializes the Lobula object
+        """
+        super().__init__()
+        self.a = 0  # Parameter a
+        self.b = 0  # Parameter b
+        self.c = 1  # Parameter c
+
+    def setup(self):
+        """
+        Initialization method
+        """
+        pass
+
+    def forward(self, onSignal, offSignal):
+        """
+        Processing method
+        Processes the input ON and OFF signals
+        
+        Parameters:
+        - varagein: Tuple containing ON and OFF signals
+        
+        Returns:
+        - lobulaOpt: Output of the Lobula layer
+        """
+        
+        # Compute Lobula output using the provided formula
+        self.output = self.a*onSignal + self.b*offSignal + self.c*onSignal*offSignal
+        
+        return self.output
+
+
+class Mi1(GammaDelay):
+    pass
+
+
+class Tm1(GammaDelay):
+    pass
+
+
+class Tm2(BaseCore):
+    """
+    Tm2 Medulla Layer Neurons in ESTMD
+    """
+   
+    def __init__(self, device='cpu'):
+        """
+        Constructor method
+        Initializes the Tm2 object
+        """
+        super().__init__()
+        self.spatial_inhibition = SpatialInhibition()  # Initialize SurroundInhibition object
+
+    def setup(self):
+        """
+        Initialization method
+        Initializes the SurroundInhibition object
+        """
+        self.spatial_inhibition.setup()
+
+    def forward(self, x):
+        """
+        Processing method
+        Applies the Surround Inhibition mechanism to the input matrix iptMatrix
+        
+        Parameters:
+        - iptMatrix: Input matrix
+        
+        Returns:
+        - tm2Opt: Output of the Tm2 layer
+        """
+        # Extract the OFF signal from iptMatrix
+        L_OFF = torch.clamp(-x, min=0)
+
+        # Process the OFF signal using SurroundInhibition
+        self.output = self.spatial_inhibition.forward(L_OFF)  
+
+        return self.output
+
+
+class Tm3(BaseCore):
+    """ Tm3  """
+    
+    def __init__(self, device='cpu'):
+        """ Constructor method
+
+        Initializes the Tm3 object
+        """
+        super().__init__()
+        self.spatial_inhibition = SpatialInhibition()  # Initialize SurroundInhibition object
+
+    def setup(self):
+        """ Initialization method
+
+        Initializes the SurroundInhibition object
+        """
+        self.spatial_inhibition.setup()
+
+    def forward(self, iptMatrix):
+        """ Processing method
+        
+        Description:
+            Applies Surround Inhibition to the On-signal matrix iptMatrix
+        
+        Parameters:
+        - iptMatrix: Input matrix
+        
+        Returns:
+        - tm3Opt: Output of the Tm3 layer
+        """
+
+        L_ON = torch.clamp(iptMatrix, min=0)
+
+        self.output = self.spatial_inhibition.forward(L_ON)  # Processes the On-signal using SurroundInhibition
+  
+        return self.output
+
+
+class LaminaLateralInhibition(BaseCore):
+    """ LAMINALATERALINHIBITION Lateral inhibition in the Lamina layer
+    
+    This class implements the lateral inhibition mechanism in the Lamina layer
+    of the ESTMD using pure PyTorch operations.
+
+    References:
+    * S. D. Wiederman, P. A. Shoemarker, D. C. O'Carroll, A model
+    for the detection of moving targets in visual clutter inspired by
+    insect physiology, PLoS ONE 3 (7) (2008) e2784.
+    * Wang H, Peng J, Yue S. A directionally selective small target
+    motion detecting visual neural network in cluttered backgrounds[J].
+    IEEE transactions on cybernetics, 2018, 50(4): 1541-1555.
+    """
+
+    def __init__(self, 
+                 sizeW1=[11, 11, 7], 
+                 lambda1=3.0, 
+                 lambda2=9.0, 
+                 sigma1=1.5, 
+                 sigma2=None):
+        """
+        Constructor
+        Initializes the LaminaLateralInhibition module
+        """
+        super().__init__()
+        self.sizeW1 = sizeW1
+        self.lambda1 = lambda1
+        self.lambda2 = lambda2
+        self.sigma1 = sigma1
+        self.sigma2 = 2.0 * sigma1 if sigma2 is None else sigma2
+        self.T = sizeW1[2]  # Temporal length
+
+        self.register_buffer('spatial_pos_kernel', torch.empty(0))
+        self.register_buffer('spatial_neg_kernel', torch.empty(0))
+        self.register_buffer('temporal_pos_kernel', torch.empty(0))
+        self.register_buffer('temporal_neg_kernel', torch.empty(0))
+
+        self.setup()  # Initialize kernels and buffers
+
+    def _create_spatial_kernels(self):
+        """初始化 DoG (Difference of Gaussian) 空间感受野权重"""
+        g_sigma2 = create_2d_gaussian_kernel(self.sizeW1[:2], self.sigma1)
+        g_sigma3 = create_2d_gaussian_kernel(self.sizeW1[:2], self.sigma2)
+        diff_of_gaussian = g_sigma2 - g_sigma3
+
+        # W_{S}^{P} 和 W_{S}^{N}
+        pos_kernel = torch.clamp(diff_of_gaussian, min=0)
+        neg_kernel = torch.clamp(diff_of_gaussian, max=0)
+
+        # 调整形状为 (out_channels=1, in_channels=1, H, W) 以匹配 F.conv2d 的需求
+        pos_kernel = pos_kernel.view(1, 1, *self.sizeW1[:2])
+        neg_kernel = neg_kernel.view(1, 1, *self.sizeW1[:2])
+        
+        return pos_kernel, neg_kernel
+
+    def _create_temporal_kernels(self):
+        """初始化时间衰减权重"""
+        t = torch.arange(self.T, dtype=torch.float32)
+        
+        # W_{T}^{P} 和 W_{T}^{N}
+        w_t_pos = torch.exp(-t / self.lambda1) / self.lambda1
+        w_t_neg = torch.exp(-t / self.lambda2) / self.lambda2
+        
+        # 调整形状为 (T, 1, 1, 1) 方便后续与 (T, Batch, C, H, W) 张量进行广播乘法
+        return w_t_pos.view(-1, 1, 1, 1), w_t_neg.view(-1, 1, 1, 1)
+
+    def setup(self):
+        # 1. 预计算空间卷积核 (Spatial Kernels)
+        spatial_pos, spatial_neg = self._create_spatial_kernels()
+        # 使用 register_buffer，这样模型调用 .cuda() 或 .to(device) 时，核也会自动转移
+        self.spatial_pos_kernel.data = spatial_pos
+        self.spatial_neg_kernel.data = spatial_neg
+
+        # 2. 预计算时间卷积核 (Temporal Kernels)
+        temporal_pos, temporal_neg = self._create_temporal_kernels()
+        self.temporal_pos_kernel.data = temporal_pos
+        self.temporal_neg_kernel.data = temporal_neg
+
+        # 3. 初始化时序状态缓存区 (采用双端队列 deque 实现高效的滑动窗口)
+        self.pos_buffer = deque(maxlen=self.T)
+        self.neg_buffer = deque(maxlen=self.T)
+
+    def reset_buffer(self):
+        """重置时序状态缓存区"""
+        self.pos_buffer.clear()
+        self.neg_buffer.clear()
+
+    def forward(self, x):
+        """
+        输入: 
+            x: 形状为 (B, C, H, W) 或 (H, W) 的张量
+        """
+
+        # === 1. 空间侧抑制 (Spatial Lateral Inhibition) ===
+        on_conv = F.conv2d(x, self.spatial_pos_kernel, padding='same')
+        off_conv = F.conv2d(x, self.spatial_neg_kernel, padding='same')
+
+        # 记录当前帧结果到时序缓存区 (新的帧在队列右侧)
+        self.pos_buffer.append(on_conv)
+        self.neg_buffer.append(off_conv)
+
+        # === 2. 时序卷积 (Temporal Convolution) ===
+
+        pos_out = compute_temporal_conv_inplace(self.pos_buffer, self.temporal_pos_kernel)
+        neg_out = compute_temporal_conv_inplace(self.neg_buffer, self.temporal_neg_kernel)
+
+        return pos_out + neg_out
+
+    
+
+
+
+
+