File size: 3,590 Bytes
a10ba7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import torch
import torch.nn as nn
from torchvision import models
class L2CS(nn.Module):
def __init__(self, backbone_name='resnet50', num_bins=90, pretrained=True):
"""
L2CS-Net Teacher Model Wrapper.
Args:
backbone_name (str): ResNet variant to use ('resnet18', 'resnet34', 'resnet50').
num_bins (int): Number of bins for pitch and yaw classification (default 90).
pretrained (bool): Whether to use ImageNet weights for the backbone.
"""
super(L2CS, self).__init__()
# 1. Initialize Backbone
if backbone_name == 'resnet18':
self.backbone = models.resnet18(weights='IMAGENET1K_V1' if pretrained else None)
feat_dim = 512
elif backbone_name == 'resnet34':
self.backbone = models.resnet34(weights='IMAGENET1K_V1' if pretrained else None)
feat_dim = 512
else: # Default resnet50
self.backbone = models.resnet50(weights='IMAGENET1K_V2' if pretrained else None)
feat_dim = 2048
# Remove the original FC layer
self.backbone = nn.Sequential(*(list(self.backbone.children())[:-1]))
# 2. Dual-Branch Heads (Pitch & Yaw)
self.fc_pitch = nn.Linear(feat_dim, num_bins)
self.fc_yaw = nn.Linear(feat_dim, num_bins)
# 3. Parameters for expectation calculation
self.num_bins = num_bins
# Bin range is typically -90 to 90 degrees or 0 to 180.
# L2CS usually uses -90 to 90 for a total of 180 degrees.
# We'll create a register_buffer for the idx tensor (0 to num_bins-1)
self.register_buffer('idx_tensor', torch.arange(num_bins).float())
def forward(self, x):
"""
Forward pass.
Returns:
pitch_logits, yaw_logits: Raw scores for each bin.
"""
# Feature extraction
features = self.backbone(x)
features = features.view(features.size(0), -1)
# Predict bins
pitch_logits = self.fc_pitch(features)
yaw_logits = self.fc_yaw(features)
return pitch_logits, yaw_logits
def get_angles(self, pitch_logits, yaw_logits):
"""
Convert logits to continuous angles (degrees) using soft-argmax (expectation).
Assumes bins represent -90 to 90 degrees with a 2-degree step for 90 bins.
"""
pitch_softmax = nn.functional.softmax(pitch_logits, dim=1)
yaw_softmax = nn.functional.softmax(yaw_logits, dim=1)
# Expectation: sum(prob * idx)
pitch_idx = torch.sum(pitch_softmax * self.idx_tensor, dim=1)
yaw_idx = torch.sum(yaw_softmax * self.idx_tensor, dim=1)
# Map bin index back to degrees: (idx * 2) - 90
# (Assuming 90 bins covering 180 degrees)
pitch_deg = pitch_idx * 2 - 90
yaw_deg = yaw_idx * 2 - 90
return pitch_deg, yaw_deg
def load_teacher_model(path=None, backbone='resnet50', num_bins=90, device='cpu'):
"""
Utility to load the teacher model with pre-trained weights.
"""
model = L2CS(backbone_name=backbone, num_bins=num_bins, pretrained=True)
if path:
# Load custom weights if provided (from L2CS-Net repo)
checkpoint = torch.load(path, map_location=device)
# Handle cases where the state_dict is nested
state_dict = checkpoint.get('state_dict', checkpoint)
model.load_state_dict(state_dict, strict=False)
model.to(device)
model.eval()
return model
|