xttmp 2.3.0__py3-none-any.whl

xttmp/core/stfeedbackstmd_core.py

@@ -0,0 +1,233 @@
+from .base_core import BaseCore
+from .math_operator import SpatialInhibition, GammaDelay
+from . import estmd_backbone 
+from ..util.create_kernel import *
+from ..util.compute_module import slice_matrix_holding_size
+
+
+class Medulla(estmd_backbone.Medulla):
+    # Medulla layer of the motion detection system
+
+    def __init__(self):
+        # Constructor method
+        # Initializes the Medulla object
+        super().__init__()
+
+        self.hPara5Mi1 = None
+        self.hPara5Tm1 = None
+        self.cellTm1Ipt = None
+
+    def setup(self):
+        # Initialization method
+        # This method initializes the Medulla layer components
+        super().setup()
+
+        self.hTm1 = GammaDelay(5, 25)
+        self.hPara5Mi1 = GammaDelay(25, 30)
+        self.hPara5Tm1 = GammaDelay(25, 30)
+
+        self.cellTm1Ipt = CircularList()
+
+        self.hTm1.setup(False)
+        self.hPara5Mi1.setup()
+        self.hPara5Tm1.setup(False)
+
+        if not self.cellTm1Ipt.initLen:
+            self.cellTm1Ipt.initLen = max(self.hPara5Mi1.lenKernel, self.hPara5Tm1.lenKernel)
+
+        self.cellTm1Ipt.reset()
+
+    def forward(self, MedullaIpt):
+        # Processing method
+        # Applies processing to the input and returns the output
+        # Process Tm2 and Tm3 components
+        tm2Signal = self.hTm2.forward(MedullaIpt)
+        tm3Signal = self.hTm3.forward(MedullaIpt)
+
+        # Process Tm1 component using output of Tm2
+        self.cellTm1Ipt.record_next(tm3Signal)
+        tm1Para3Signal = self.hTm1.forward(self.cellTm1Ipt)
+        tm1Para5Signal = self.hPara5Tm1.forward(self.cellTm1Ipt)
+
+        # Process Mi1 component using output of Tm3
+        mi1Para5Signal = self.hPara5Mi1.forward(tm3Signal)
+
+        # Store the output signals in Opt property
+        varageout = [tm3Signal, tm1Para3Signal, mi1Para5Signal, tm2Signal, tm1Para5Signal, self.hPara5Mi1.tau]
+        self.Opt = varageout
+        return varageout
+
+
+class Lobula(BaseCore):
+    # Lobula layer of the motion detection system
+
+    def __init__(self):
+        # Constructor method
+        # Initializes the Lobula object
+        super().__init__()
+
+        self.hSTMD = None
+        self.hLPTC = None
+
+    def setup(self):
+        # Initialization method
+        # This method initializes the Lobula layer component
+        self.hSTMD = Stmdcell()
+        self.hLPTC = Lptcell()
+        self.hSTMD.setup()
+        self.hLPTC.setup(self.hSTMD.hGammaDelay.lenKernel)
+
+    def forward(self, varagein):
+        # Processing method
+        # Performs a correlation operation on the ON and OFF channels
+        # and then applies surround inhibition
+
+        # Extract ON and OFF channel signals from the input
+        tm3Signal, tm1Para3Signal, mi1Para5Signal, tm2Signal, tm1Para5Signal, tau5 = varagein
+
+        psi, fai = self.hLPTC.forward(tm3Signal, mi1Para5Signal, tm2Signal, tm1Para5Signal, tau5)
+
+        lobulaOpt = self.hSTMD.forward(tm3Signal, tm1Para3Signal, psi, fai)
+
+        # Store the output in Opt property
+        self.Opt = lobulaOpt
+
+        return lobulaOpt, []
+
+
+class Stmdcell(BaseCore):
+    # Lobula layer of the motion detection system
+
+    def __init__(self):
+        # Constructor method
+        # Initializes the Lobula object
+        super().__init__()
+
+        self.hSubInhi = None  # SpatialInhibition component
+        self.alpha = 0.1  # Parameter alpha
+        self.gaussKernel = None  # Gaussian kernel
+        self.hGammaDelay = None
+        self.cellDPlusE = None
+        self.paraGaussKernel = {'size': 3, 'eta': 1.5}
+
+    def setup(self):
+        # Initialization method
+        # This method initializes the Lobula layer component
+        self.hSubInhi = SpatialInhibition()
+        self.hGammaDelay = GammaDelay(6, 12)
+        self.cellDPlusE = CircularList()
+
+        self.hSubInhi.setup()
+        self.hGammaDelay.setup()
+
+        if not self.cellDPlusE.initLen:
+            self.cellDPlusE.initLen = self.hGammaDelay.lenKernel
+        self.cellDPlusE.reset()
+
+        self.gaussKernel = gaussian_filter(
+            np.zeros((self.paraGaussKernel['size'], self.paraGaussKernel['size'])),
+            self.paraGaussKernel['eta']
+        )
+
+    def forward(self, tm3Signal, tm1Signal, faiList, psiList):
+        # Processing method
+        # Performs temporal convolution, correlation, and surround inhibition
+        convnIpt = [None] * self.cellDPlusE.initLen
+
+        for idxT in range(len(convnIpt)-1, -1, -1):
+            pointer = self.cellDPlusE.pointer
+            if self.cellDPlusE[pointer] is not None:
+                fai = faiList[idxT]
+                psi = psiList[idxT]
+                convnIpt[idxT] = slice_matrix_holding_size(self.cellDPlusE[pointer], psi, fai)
+                pointer = (pointer - 1) % self.cellDPlusE.initLen
+
+        feedbackSignal = self.hGammaDelay.forward_list(convnIpt)         
+
+        if feedbackSignal is not None:
+            feedbackSignal *= self.alpha 
+            correlationD = np.maximum(tm3Signal - feedbackSignal, 0) * np.maximum(tm1Signal - feedbackSignal, 0)
+        else:
+            correlationD = np.maximum(tm3Signal, 0) * np.maximum(tm1Signal, 0)
+
+        correlationE = filter2D(tm3Signal * tm1Signal, -1, self.gaussKernel, borderType=BORDER_CONSTANT)
+
+        lateralInhiSTMDOpt = self.hSubInhi.forward(correlationD)
+
+        self.cellDPlusE.record_next(correlationD + correlationE)
+
+        self.Opt = lateralInhiSTMDOpt
+
+        return lateralInhiSTMDOpt
+
+
+class Lptcell(BaseCore):
+    # Lptcell Lobula Plate Tangential Cell
+
+    def __init__(self):
+        # Constructor method
+        # Initializes the Lobula object
+        super().__init__()
+
+        self.bataList = list(range(2, 19, 2))
+        self.thetaList = np.arange(0, 2 * np.pi, np.pi / 4)
+        self.velocity = None
+        self.tuningCurvef = None
+
+    def setup(self, lenVelocity):
+        self.velocity = np.zeros(lenVelocity)
+
+        lenBataList = len(self.bataList)
+        lenThetaList = len(self.thetaList)
+        # generate gauss distribution
+        gaussianDistribution = np.exp(-0.5 * ((np.arange(-199, 201) - 1) / (100 / 2)) ** 2)
+        # normalization
+        gaussianDistribution /= np.max(gaussianDistribution)
+
+        self.tuningCurvef = np.zeros((lenBataList, lenThetaList * 100 + 200))
+        self.tuningCurvef[0, :300] = gaussianDistribution[100:400]
+        self.tuningCurvef[-1, -300:] = gaussianDistribution[:300]
+        for id in range(1, lenBataList - 1):
+            idRange = slice((id+1) * 100 - 200, (id+1) * 100 + 200)
+            self.tuningCurvef[id, idRange] = gaussianDistribution
+
+    def forward(self, tm1Signal, tm2Signal, tm3Signal, mi1Signal, tau5):
+        lenBataList = len(self.bataList)
+        lenThetaList = len(self.thetaList)
+        sumLplcOptR = np.zeros((lenBataList, lenThetaList))
+
+        for idBata, bata in enumerate(self.bataList):
+            for idTheta, theta in enumerate(self.thetaList):
+                shiftX = np.round(bata * np.cos(theta + np.pi / 2)).astype(int)
+                shiftY = np.round(bata * np.sin(theta + np.pi / 2)).astype(int)
+                shiftMi1Signal = slice_matrix_holding_size(mi1Signal, shiftY, shiftX)
+                shiftTm1Signal = slice_matrix_holding_size(tm1Signal, shiftY, shiftX)
+
+                ltlcOpt = tm3Signal * shiftMi1Signal + tm2Signal * shiftTm1Signal
+
+                sumLplcOpt = np.sum(ltlcOpt)
+
+                sumLplcOptR[idBata, idTheta] = sumLplcOpt
+
+        # preferTheta
+        firingRate = np.max(sumLplcOptR, axis=1)
+        preferTheta = np.argmax(sumLplcOptR, axis=1)
+        maxTheta = np.max(preferTheta)
+
+        # background velocity
+        self.velocity = np.roll(self.velocity, -1)
+        temp = [np.sum((firingRate[i] - self.tuningCurvef[i,:]) ** 2) for i in range(len(firingRate))]
+        self.velocity[-1] = np.argmin(temp)
+
+        sumV = np.zeros_like(self.velocity)
+        for idV in range(len(self.velocity)):
+            sumV[idV] = np.sum(self.velocity[idV:])
+
+        fai = sumV * np.cos(maxTheta)
+        psi = sumV * np.sin(maxTheta)
+
+        self.Opt = [fai, psi]
+
+        return fai, psi
+
+

xttmp/core/stmdplus_core.py

@@ -0,0 +1,187 @@
+import torch
+from torch.nn import functional as F
+
+from .base_core import BaseCore
+from ..util.create_kernel import create_T1_kernels
+from ..util.compute_module import AreaNMS
+from ..util.compute_module import compute_response
+
+
+class ContrastPathway(BaseCore):
+    """ContrastPathway class for ApgSTMD - PyTorch Version."""
+
+    def __init__(self):
+        """Constructor method."""
+        super().__init__()
+        
+        self.theta = torch.tensor([0, torch.pi/4, torch.pi/2, 3*torch.pi/4])
+        self.alpha2 = 1.5
+        self.eta = 3
+        self.sizeT1 = 11
+        
+        self.register_buffer('T1_kernel', torch.empty(0))
+
+    def setup(self):
+        """Initialization method."""
+        # 假设 create_T1_kernels 返回的是 4 个 kernel 的 List 或 NumPy array
+        _T1_kernel = create_T1_kernels(len(self.theta), self.alpha2, self.eta, self.sizeT1)
+        self.T1_kernel.data = _T1_kernel
+
+    def forward(self, x):
+        """
+        Processing method.
+        retinaOpt: 可以是 (H, W) 或 (1, 1, H, W) 的 PyTorch Tensor
+        """
+        
+        # out 的形状是 (1, 4, H, W)
+        self.Opt = F.conv2d(x, self.T1_kernel, padding='same')
+
+        return self.Opt
+    
+
+class MushroomBody(BaseCore):
+    # MushroomBody class for STMDPlus - Ultimate Vectorized PyTorch Version
+
+    def __init__(self):
+        super().__init__()
+        self.nms_size = 5
+
+        self.DBSCANDist = 5.0
+        self.lenDBSCAN = 100
+        self.SDThres = 5.0
+        
+        self.torch_nms = None
+        
+        # ================= 终极优化：预分配的全张量状态机 =================
+        self.C = None           # 通道数 (将在第一帧自动推断)
+        self.trackID = None     # 张量: [N, D] 坐标
+        self.trackInfo = None   # 张量: [N, C, lenDBSCAN] (固定大小环形缓冲区)
+        self.trackLens = None   # 张量: [N] (记录当前轨迹有效长度)
+        self.trackPtr = None    # 张量: [N] (记录环形缓冲区的写入指针)
+
+    def setup(self):
+        self.torch_nms = AreaNMS(self.nms_size)
+
+    def forward(self, lobulaOpt, contrast_tensor):
+        device = lobulaOpt.device
+
+        maxLobulaOpt = compute_response(lobulaOpt) 
+        nmsLobulaOpt = self.torch_nms(maxLobulaOpt)
+
+        mask_not_nms = (nmsLobulaOpt == 0)
+        mushroomBodyOpt = lobulaOpt * mask_not_nms
+
+        maxNumber = torch.max(nmsLobulaOpt)
+        if maxNumber <= 0:
+            self.trackID = None
+            return mushroomBodyOpt
+
+        # --- 获取新检测点 ---
+        newID = torch.nonzero(nmsLobulaOpt > 0).float() 
+        if len(newID) == 0:
+            self.trackID = None
+            return mushroomBodyOpt
+
+        curr_y, curr_x = newID[:, -2].long(), newID[:, -1].long()
+        # all_new_contrasts 形状 [C, M], M是新目标数
+        all_new_contrasts = contrast_tensor[:, 0, curr_y, curr_x]
+        
+        if self.C is None:
+            self.C = contrast_tensor.shape[0]
+
+        M = len(newID)
+
+        # ================= 1. 轨迹匹配 (保持 CPU 高效碰撞处理) =================
+        matched_old, matched_new, used_new = [], [], set()
+        
+        if self.trackID is not None and len(self.trackID) > 0:
+            DD = torch.cdist(self.trackID[:, -2:], newID[:, -2:])
+            D1, min_idx = torch.min(DD, dim=1)
+            
+            # 转移到 CPU 做极速冲突判定
+            D1_cpu, min_idx_cpu = D1.cpu().numpy(), min_idx.cpu().numpy()
+            for i, d1 in enumerate(D1_cpu):
+                if d1 <= self.DBSCANDist:
+                    j = min_idx_cpu[i]
+                    if j not in used_new:
+                        used_new.add(j)
+                        matched_old.append(i)
+                        matched_new.append(j)
+
+        # ================= 2. 环形缓冲区批量状态更新 (0 For循环) =================
+        matched_old_ts = torch.tensor(matched_old, dtype=torch.long, device=device)
+        matched_new_ts = torch.tensor(matched_new, dtype=torch.long, device=device)
+        unmatched_new_ts = torch.tensor([j for j in range(M) if j not in used_new], dtype=torch.long, device=device)
+
+        if self.trackID is not None and len(matched_old) > 0:
+            # --- 提取续航的轨迹状态 ---
+            next_trackID = newID[matched_new_ts]
+            next_trackInfo = self.trackInfo[matched_old_ts]
+            next_trackLens = self.trackLens[matched_old_ts]
+            next_trackPtr = self.trackPtr[matched_old_ts]
+            
+            # 批量写入环形缓冲区 (全矩阵操作)
+            batch_idx = torch.arange(len(matched_old), device=device)
+            next_trackInfo[batch_idx, :, next_trackPtr] = all_new_contrasts[:, matched_new_ts].T
+            next_trackLens = torch.clamp(next_trackLens + 1, max=self.lenDBSCAN)
+            next_trackPtr = (next_trackPtr + 1) % self.lenDBSCAN
+        else:
+            # 定义空张量用于拼接
+            next_trackID = torch.empty((0, newID.shape[1]), device=device)
+            next_trackInfo = torch.empty((0, self.C, self.lenDBSCAN), device=device)
+            next_trackLens = torch.empty((0,), dtype=torch.long, device=device)
+            next_trackPtr = torch.empty((0,), dtype=torch.long, device=device)
+
+        if len(unmatched_new_ts) > 0:
+            # --- 批量初始化新轨迹 ---
+            N_add = len(unmatched_new_ts)
+            add_trackID = newID[unmatched_new_ts]
+            add_trackInfo = torch.zeros((N_add, self.C, self.lenDBSCAN), device=device)
+            add_trackInfo[:, :, 0] = all_new_contrasts[:, unmatched_new_ts].T
+            add_trackLens = torch.ones(N_add, dtype=torch.long, device=device)
+            add_trackPtr = torch.ones(N_add, dtype=torch.long, device=device)
+            
+            # --- 与续航的轨迹合并 ---
+            self.trackID = torch.cat([next_trackID, add_trackID], dim=0)
+            self.trackInfo = torch.cat([next_trackInfo, add_trackInfo], dim=0)
+            self.trackLens = torch.cat([next_trackLens, add_trackLens], dim=0)
+            self.trackPtr = torch.cat([next_trackPtr, add_trackPtr], dim=0)
+        else:
+            self.trackID, self.trackInfo, self.trackLens, self.trackPtr = next_trackID, next_trackInfo, next_trackLens, next_trackPtr
+
+        # ================= 3. 并行 STD 计算与批量擦除 =================
+        if self.trackID is not None and len(self.trackID) > 0:
+            valid_mask = self.trackLens > 1
+            if valid_mask.any():
+                # 只取出长度 > 1 的轨迹
+                chk_info = self.trackInfo[valid_mask]    # [K, C, 100]
+                chk_lens = self.trackLens[valid_mask]    # [K]
+                chk_coords = self.trackID[valid_mask]    # [K, D]
+                
+                # 创建时间遮罩: 标记环形缓冲区中哪些数据是有效的
+                time_idx = torch.arange(self.lenDBSCAN, device=device).view(1, 1, -1) # [1, 1, 100]
+                data_mask = time_idx < chk_lens.view(-1, 1, 1) # [K, 1, 100]
+                
+                # 手动并行计算 Masked STD (利用数学公式：Var = sum((x - mean)^2) / (N - 1))
+                # 因为方差是无序的，即便环形缓冲区数据没按时间排序，也绝对不影响最终计算结果！
+                sum_val = (chk_info * data_mask).sum(dim=-1) # [K, C]
+                mean_val = sum_val / chk_lens.unsqueeze(1).float()
+                
+                diff_sq = ((chk_info - mean_val.unsqueeze(-1)) * data_mask) ** 2
+                var_val = diff_sq.sum(dim=-1) / (chk_lens.unsqueeze(1).float() - 1.0)
+                std_val = var_val.sqrt() # [K, C] 
+                
+                # 找出最大的 STD，并判定
+                max_std, _ = std_val.max(dim=1) # [K]
+                erase_mask = max_std < self.SDThres # [K]
+                
+                if erase_mask.any():
+                    erase_coords = chk_coords[erase_mask]
+                    e_y = erase_coords[:, -2].long()
+                    e_x = erase_coords[:, -1].long()
+                    # 终极一键批量擦除
+                    mushroomBodyOpt[..., e_y, e_x] = 0
+
+        self.Opt = mushroomBodyOpt
+        return mushroomBodyOpt
+

xttmp/core/stmdplusv2_core.py

@@ -0,0 +1,82 @@
+import numpy as np
+from scipy.spatial.distance import cdist
+
+from . import stmdplus_core
+
+class MushroomBody(stmdplus_core.MushroomBody):
+    # MushroomBody class for STMDPlus
+
+    def __init__(self):
+        # Constructor method
+        # Initializes the MushroomBody object
+        super().__init__()
+
+
+    def setup(self):
+        # Initialization method
+        # Initializes the non-maximum suppression
+        super().setup()
+
+    def forward(self, lobulaOpt, contrastOpt):
+        # Processing method
+        # Processes the input lobulaOpt and contrastOpt to generate mushroomBodyOpt
+
+        nmsLobulaOpt = self.hNMS.nms(lobulaOpt)
+
+        mushroomBodyOpt = lobulaOpt * (nmsLobulaOpt > 0)
+
+        maxNumber = np.max(nmsLobulaOpt)
+
+        if maxNumber <= 0:
+            self.trackID = None
+            self.trackInfo = []
+            return mushroomBodyOpt
+
+        idX, idY = np.where(nmsLobulaOpt > 0)
+        newID = np.column_stack((idX, idY))
+
+        shouldTrackID = np.ones(len(self.trackID), dtype=bool) if self.trackID is not None else np.array([], dtype=bool)
+        shouldAddNewID = np.ones(len(idX), dtype=bool)
+        numContrast = len(contrastOpt)
+
+        if self.trackID is not None:
+            DD = cdist(self.trackID, newID)
+            D1 = np.min(DD, axis=1)
+
+            for idxI, d1 in enumerate(D1):
+                if d1 <= self.DBSCANDist:
+                    idxJ = np.argmin(DD[idxI])
+                    if shouldAddNewID[idxJ]:
+                        self.trackID[idxI] = newID[idxJ]
+                        nowContrast = np.array(
+                            [[contrastOpt[idCont][newID[idxJ, 0], newID[idxJ, 1]]] for idCont in range(numContrast)])
+                        self.trackInfo[idxI] = np.hstack((self.trackInfo[idxI], nowContrast))
+                        shouldTrackID[idxI] = False
+                        shouldAddNewID[idxJ] = False
+
+            self.trackID = np.delete(self.trackID, np.where(shouldTrackID), axis=0)
+            self.trackInfo = [x for idx, x in enumerate(self.trackInfo) if not shouldTrackID[idx]]
+
+        oldTractNum = len(self.trackInfo)
+
+        isxNew = np.where(shouldAddNewID)[0]
+        for kk in isxNew:
+            if self.trackID is None:
+                self.trackID = newID[kk]
+            else:
+                self.trackID = np.vstack((self.trackID, newID[kk]))
+            nowContrast = np.array(
+                [[contrastOpt[idCont][newID[kk, 0], newID[kk, 1]]] for idCont in range(numContrast)])
+            self.trackInfo.append(nowContrast)
+
+        for idx in range(oldTractNum):
+            if np.max(np.std(self.trackInfo[idx], axis=1)) < self.SDThres:
+                idX = self.trackID[idx, 0]
+                idY = self.trackID[idx, 1]
+                mushroomBodyOpt[idX, idY] = 0
+
+            if self.trackInfo[idx].shape[1] > self.lenDBSCAN:
+                self.trackInfo[idx] = self.trackInfo[idx][:, 1:]
+
+        self.Opt = mushroomBodyOpt
+        return mushroomBodyOpt