| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| class AevanisTraverseHead(nn.Module):
|
| def __init__(self, visual_dim=256, hidden_dim=128):
|
| super().__init__()
|
| self.visual_dim = visual_dim
|
|
|
|
|
| self.dense_conv = nn.Sequential(
|
| nn.ConvTranspose2d(visual_dim, hidden_dim, kernel_size=4, stride=2, padding=1),
|
| nn.BatchNorm2d(hidden_dim),
|
| nn.GELU(),
|
| nn.Conv2d(hidden_dim, 32, kernel_size=3, padding=1),
|
| nn.GELU(),
|
| nn.Conv2d(32, 1, kernel_size=1),
|
| nn.Sigmoid()
|
| )
|
|
|
| self.global_pool = nn.AdaptiveAvgPool1d(1)
|
|
|
|
|
| self.slope_head = nn.Sequential(
|
| nn.Linear(visual_dim, hidden_dim),
|
| nn.LayerNorm(hidden_dim),
|
| nn.GELU(),
|
| nn.Linear(hidden_dim, 2)
|
| )
|
|
|
|
|
| self.slip_head = nn.Sequential(
|
| nn.Linear(visual_dim, hidden_dim),
|
| nn.LayerNorm(hidden_dim),
|
| nn.GELU(),
|
| nn.Linear(hidden_dim, 1),
|
| nn.Sigmoid()
|
| )
|
|
|
| def forward(self, visual_tokens):
|
| B, N, C = visual_tokens.shape
|
| h = w = int(N ** 0.5)
|
| feat_2d = visual_tokens.permute(0, 2, 1).reshape(B, C, h, w)
|
|
|
| foothold_map = self.dense_conv(feat_2d).squeeze(1)
|
| global_feat = self.global_pool(visual_tokens.permute(0, 2, 1)).squeeze(-1)
|
|
|
| slope_deg = self.slope_head(global_feat) * 45.0
|
| slip_risk = self.slip_head(global_feat)
|
|
|
| return {
|
| "foothold_map": foothold_map,
|
| "slope_angles": slope_deg,
|
| "slip_risk": slip_risk
|
| }
|
|
|
|
|
| class AevanisTraverseAgent(nn.Module):
|
| def __init__(self, model_scale="nano", head_weights=None, backbone=None):
|
| super().__init__()
|
| self.model_scale = model_scale
|
| visual_dim = 256 if model_scale == "nano" else 384
|
| self.traverse_head = AevanisTraverseHead(visual_dim=visual_dim)
|
| self.backbone = backbone
|
|
|
| if head_weights is not None:
|
| self.load_head(head_weights)
|
|
|
| def load_head(self, path: str):
|
| state = torch.load(path, map_location="cpu")
|
| cleaned = {k.replace("traverse_head.", ""): v for k, v in state.items()}
|
| self.traverse_head.load_state_dict(cleaned, strict=False)
|
| self.eval()
|
|
|
| def forward(self, visual_input):
|
| if visual_input.dim() == 4:
|
| if self.backbone is not None:
|
| with torch.no_grad():
|
| features = self.backbone(visual_input)
|
| tokens = features["spatial_tokens"]
|
| else:
|
| B = visual_input.shape[0]
|
| tokens = torch.zeros(B, 196, self.traverse_head.visual_dim)
|
| else:
|
| tokens = visual_input
|
|
|
| return self.traverse_head(tokens)
|
|
|
| @classmethod
|
| def from_pretrained(cls, model_scale="nano", head_weights=None):
|
| return cls(model_scale=model_scale, head_weights=head_weights)
|
|
|
