Radar

Configs

radar_configs = {
    'c': 3e8, 
    'startFreq': 77e9, 
    'Tr': 60e-6,                    # Sweeping frequency time
    'Idle_time': 100e-6,            # free time
    'Fs': 10e6,                     # Sampling frequency
    'Slope': 29.982e12,             # chirp slope
    'Bandwidth': 60e-6 * 29.982e12, # Transmission signal bandwidth
    'BandwidthValid': 0.767539200e9,# Effective bandwidth of the transmitted signal
    'range_size': 256,              # range size
    'azimuth_size': 181,            # azimuth size
    'elevation_size': 181,          # elevation size
    'crop_num': 3,                  # crop some indices in range domain
    'n_chirps': 128,                # number of chirps in one frame
    'min_azimuth': -90,             # min radar azimuth
    'max_azimuth': 90,              # max radar azimuth
    'min_elevation': -90,           # min radar elevation
    'max_elevation': 90,            # max radar elevation    
    'min_range': 1.0,               # min radar range
    'max_range': 25.0,              # max radar range
    'range_res': 3e8/(2*0.767539200e9), 
    'angle_res': 1
}

dimssnet_configs = {
    'n_epoch': 10,
    'batch_size': 2,
    'learning_rate': 1e-5,
    'range_size': 50,
    'azimuth_size': 181,
    'elevation_size': 181,
    'doppler_size': 181, 
    'min_azimuth': -90,             # min radar azimuth
    'max_azimuth': 90,              # max radar azimuth
    'min_elevation': -90,           # min radar elevation
    'max_elevation': 90,            # max radar elevation    
    'min_range': 1.0,               # min radar range
    'max_range': 25.0,              # max radar range
}

n_class = 5
class_table = {
    -1: 'background', 
    0: 'suicase', 
    1: 'chair', 
    2: 'person',
    3: 'wall'
}

Models

import torch
import torch.nn as nn
import torch.nn.functional as F

class BaseEncoder(nn.Module):
    def __init__(self):
        super(BaseEncoder, self).__init__()
        self.conv1a = nn.Conv2d(in_channels=1, out_channels=64, 
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv1b = nn.Conv2d(in_channels=64, out_channels=64, 
                                kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.conv2a = nn.Conv2d(in_channels=64, out_channels=128, 
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv2b = nn.Conv2d(in_channels=128, out_channels=128, 
                                kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.conv3a = nn.Conv2d(in_channels=128, out_channels=256, 
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv3b = nn.Conv2d(in_channels=256, out_channels=256, 
                                kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))

        self.bn1a = nn.BatchNorm2d(num_features=64)
        self.bn1b = nn.BatchNorm2d(num_features=64)
        self.bn2a = nn.BatchNorm2d(num_features=128)
        self.bn2b = nn.BatchNorm2d(num_features=128)
        self.bn3a = nn.BatchNorm2d(num_features=256)
        self.bn3b = nn.BatchNorm2d(num_features=256)

        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.bn1a(self.conv1a(x)))
        x = self.relu(self.bn1b(self.conv1b(x)))
        x = self.relu(self.bn2a(self.conv2a(x)))
        x = self.relu(self.bn2b(self.conv2b(x)))
        x = self.relu(self.bn3a(self.conv3a(x)))
        x = self.relu(self.bn3b(self.conv3b(x)))
        
        return x


class BaseDecoder(nn.Module):
    def __init__(self):
        super(BaseDecoder, self).__init__()
        self.convt1 = nn.ConvTranspose2d(in_channels=256, out_channels=128,
                                         kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.convt2 = nn.ConvTranspose2d(in_channels=128, out_channels=64,
                                         kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.convt3 = nn.ConvTranspose2d(in_channels=64, out_channels=32,
                                         kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        self.prelu = nn.PReLU()

    def forward(self, x):
        x = self.prelu(self.convt1(x))
        x = self.prelu(self.convt2(x))
        x = self.prelu(self.convt3(x))

        return x

range_size=radar_configs['range_size']
azimuth_size=radar_configs['azimuth_size']
elevation_size=radar_configs['elevation_size']

class Fuse_fea(nn.Module):
    def __init__(self):
        super(Fuse_fea, self).__init__()
        self.convt1 = nn.Conv2d(in_channels=192, out_channels=96,
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.convt2 = nn.Conv2d(in_channels=96, out_channels=48,
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.convt3 = nn.Conv2d(in_channels=48, out_channels=24,
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.convt4 = nn.Conv2d(in_channels=24, out_channels=n_class,
                                kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))        
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, feas_sra, feas_dra, feas_sre, feas_dre, feas_srd, feas_drd):
        feas_sre = torch.sum(feas_sre, 3, keepdim=True) # (B, 32, 50, 181) -> (B, 32, 50, 1)
        feas_sre = feas_sre.expand(-1, -1, -1, azimuth_size)
        
        feas_dre = torch.sum(feas_dre, 3, keepdim=True)  # (B, 32, 50, 181) -> (B, 32, 50, 1)
        feas_dre = feas_dre.expand(-1, -1, -1, azimuth_size)

        feas_srd = torch.sum(feas_srd, 3, keepdim=True)  # (B, 32, 50, 181) -> (B, 32, 50, 1)
        feas_srd = feas_srd.expand(-1, -1, -1, azimuth_size)

        feas_drd = torch.sum(feas_drd, 3, keepdim=True)  # (B, 32, 50, 181) -> (B, 32, 50, 1)
        feas_drd = feas_drd.expand(-1, -1, -1, azimuth_size)

        # Resize using bilinear interpolation
        feas_sra = F.interpolate(feas_sra, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        feas_dra = F.interpolate(feas_dra, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        feas_sre = F.interpolate(feas_sre, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        feas_dre = F.interpolate(feas_dre, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        feas_srd = F.interpolate(feas_srd, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        feas_drd = F.interpolate(feas_drd, size=(range_size, azimuth_size), mode='bilinear', align_corners=False)
        
        feas_ra = torch.cat((feas_sra, feas_dra, feas_sre, feas_dre, feas_srd, feas_drd), 1) # 6*(B, 32, 50, 181) -> (B, 192, 50, 181)

        x = self.relu(self.convt1(feas_ra))  # (B, 192, 50, 181) -> (B, 96, 50, 181)
        x = self.relu(self.convt2(x))  # (B, 96, 50, 181) -> (B, 48, 50, 181)
        x = self.relu(self.convt3(x))  # (B, 48, 50, 181) -> (B, 24, 50, 181)        
        x = self.sigmoid(self.convt4(x))  # (B, 24, 50, 181) -> (B, 5, 50, 181)
        
        return x

class DIMSSNet(nn.Module):
    def __init__(self):
        super(DIMSSNet, self).__init__()
        self.encode_sra, self.decode_sra = BaseEncoder(), BaseDecoder()
        self.encode_dra, self.decode_dra = BaseEncoder(), BaseDecoder()
        self.encode_sre, self.decode_sre = BaseEncoder(), BaseDecoder()
        self.encode_dre, self.decode_dre = BaseEncoder(), BaseDecoder()
        self.encode_srd, self.decode_srd = BaseEncoder(), BaseDecoder()
        self.encode_drd, self.decode_drd = BaseEncoder(), BaseDecoder()
        
        self.fuse_fea = Fuse_fea()
        
    def forward(self, x_sra, x_dra, x_sre, x_dre, x_srd, x_drd):
        x_sra = self.encode_sra(x_sra)
        feas_sra = self.decode_sra(x_sra)  # (B, 32, 50, 181)
        
        x_dra = self.encode_dra(x_dra)
        feas_dra = self.decode_dra(x_dra)  # (B, 32, 50, 181)
        
        x_sre = self.encode_sre(x_sre)
        feas_sre = self.decode_sre(x_sre)  # (B, 32, 50, 181)
        
        x_dre = self.encode_dre(x_dre)
        feas_dre = self.decode_dre(x_dre)  # (B, 32, 50, 181)

        x_srd = self.encode_sre(x_srd)
        feas_srd = self.decode_sre(x_srd)  # (B, 32, 50, 181)
        
        x_drd = self.encode_dre(x_drd)
        feas_drd = self.decode_dre(x_drd)  # (B, 32, 50, 181)        
        
        feas = self.fuse_fea(feas_sra, feas_dra, feas_sre, feas_dre, feas_srd, feas_drd) # (B, 32, 50, 181)

        return feas

Utils

import scipy.io
import numpy as np

# Define the target size
TARGET_HEIGHT = 50
TARGET_WIDTH = 181

def load_raw_data(file_path):
    data = scipy.io.loadmat(file_path)

    # Range-Azimuth Static
    range_az_static = data['range_az_static']
    range_az_static_prob = data['range_az_music_average_static_probability']
    range_az_static_class2_prob = range_az_static_prob[:, :, 1]

    # Range-Elevation Static
    range_ele_static = data['range_ele_static_all']

    # Range-Azimuth Dynamic
    range_az_dynamic = data['range_az_dynamic']
    range_az_dynamic_prob = data['range_az_music_average_dynamic_probability']
    range_az_dynamic_class1_prob = range_az_dynamic_prob[:, :, 0]

    # Range-Elevation Dynamic
    range_ele_dynamic = data['range_ele_dynamic_all'] 

    return [range_az_static, range_az_dynamic, range_ele_static,
            range_ele_dynamic, range_az_dynamic_class1_prob, range_az_static_class2_prob]

def generate_parameter():
    parameter = {
        'c': 3e8,  # the speed of light
        'start_freq': 77e9,  # starting frequency
        'tr': 60e-6,  # Sweeping time, that is, the cycle
        'samples': 256,  # sampling point
        'fs': 10e6,  # Sampling rate
        'tframe_set': 80e-3,  # Frame period
        'range_bin': 256,  # range bin
        'chirps': 128,  # chirp number
        'doppler_bin': 128,  # doppler bin
        'slope': 29.982e12,  # chirp slope
        'bandwidth': 29.982e12 * 60e-6,  # Effective bandwidth of the transmitted signal
        'bandwidth_valid': 256 / 10e6 * 29.982e12,  # Transmission signal bandwidth
        'center_freq': 77e9 + (29.982e12 * 60e-6) / 2,  # center frequency
        'lambda': 3e8 / (77e9 + (29.982e12 * 60e-6) / 2),  # wavelength
        'tx_antenna': [1, 1, 1],  # The number of transmitting antennas
        'rx_antenna': [1, 1, 1, 1],  # Number of receiving antennas
        'tx_num': 3,  # The number of transmitting antennas
        'rx_num': 4,  # The number of receiving antennas
        'virtual_antenna': 12,  # Number of virtual antennas
        'dz': (3e8 / (77e9 + (29.982e12 * 60e-6) / 2)) / 2,  # Pitch spacing of the receiving antenna
        'dx': (3e8 / (77e9 + (29.982e12 * 60e-6) / 2)) / 2,  # Horizontal spacing of receiving antennas
        'num_cpi': 50,  # frames
    }
    return parameter

Dataloaders

import os
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split

class CustomDataset(Dataset):
    def __init__(self, file_paths, transform=None):
        self.file_paths = file_paths
        self.transform = transform
        self.file_names = [os.path.basename(path) for path in file_paths]

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        data = torch.load(self.file_paths[idx], weights_only=True)
        
        images = data[:6]  #first 6 channels are the heatmaps
        labels = data[6]   #last channel is the label
        
        images = images.float()  
        labels = labels.long()   
        labels = labels + 1     #labels are -1 to 3, add 1 to make them 0 to 4

        if self.transform:
            images = self.transform(images)
            labels = self.transform(labels)
        
        return images, labels, self.file_names[idx]

class CustomDataset_test(Dataset):
    def __init__(self, file_paths, transform=None):
        self.file_paths = file_paths
        self.transform = transform
        self.file_names = [os.path.basename(path) for path in file_paths]

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        data = torch.load(self.file_paths[idx], weights_only=True)
        
        images = data[:6] #only load first 6 channels, labels are not in val_set and test_set
        
        images = images.float()  

        if self.transform:
            images = self.transform(images)
        
        return images, self.file_names[idx]

# extend base_path to file_path
def generate_file_paths(base_path):
    file_paths = []
    for frame in os.listdir(base_path):
        frame_path = os.path.join(base_path, frame)
        if frame_path.endswith('.mat.pt'):
            file_paths.append(frame_path)
    return [path for path in file_paths if os.path.exists(path)]

def load_data(base_path, batch_size=4, num_workers=2, test_size=0.2):
    file_paths = generate_file_paths(base_path)
    
    train_paths, test_paths = train_test_split(file_paths, test_size=test_size, random_state=42)
    
    train_dataset = CustomDataset(file_paths=train_paths)
    test_dataset = CustomDataset(file_paths=test_paths)
    
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True, 
        num_workers=num_workers, 
        drop_last=True
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=num_workers, 
        drop_last=True
    )
    
    return train_loader, test_loader

Loss Function

#Used to handle the problem of class imbalance

def to_one_hot(tensor, nClasses):
    n, h, w = tensor.size()
    one_hot = torch.zeros(n, nClasses, h, w, device=tensor.device)
    index_tensor = tensor.long().view(n, 1, h, w)
    one_hot = one_hot.scatter_(1, index_tensor, 1)
    return one_hot

def _neg_loss(pred, gt):
    pred = torch.clamp(pred, 1e-5, 1-1e-5)

    pos_inds = gt.eq(1).float()
    neg_inds = gt.eq(0).float()

    balance_cof = 1.0
    pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds * balance_cof
    neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_inds
    num_pos = pos_inds.float().sum()
    pos_loss_sum = pos_loss.sum()
    neg_loss_sum = neg_loss.sum()
    if num_pos == 0:
        loss = -neg_loss_sum
    else:
        loss = -(pos_loss_sum + neg_loss_sum) / num_pos
    return loss

class FocalLoss(nn.Module):
    def __init__(self):
        super(FocalLoss, self).__init__()
        self.neg_loss = _neg_loss

    def forward(self, out, target):
        return self.neg_loss(out, target)

Train

import torch.optim as optim
import matplotlib.pyplot as plt
import time

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

num_epochs = 20

base_path = '../training_set'
train_loader, test_loader = load_data(
    base_path=base_path,
    batch_size=4,
    num_workers=2,
    test_size=0.2
)

dimssnet = DIMSSNet().to(device)
criterion = FocalLoss()
optimizer = optim.Adam(dimssnet.parameters(), lr=0.001)
start_time = time.time()
test_loss_record = []
train_loss_record = []

for epoch in range(num_epochs):
    dimssnet.train()
    running_loss = 0.0
    for images, labels, _ in train_loader:
        images = images.to(device)
        labels = labels.to(device)
        
        images = [image.unsqueeze(1) for image in images]
        images = torch.stack(images, dim=0)
        images = images.float()
        
        outputs = dimssnet(images[:, 0, :, :], images[:, 1, :, :], images[:, 2, :, :],
                            images[:, 3, :, :], images[:, 4, :, :], images[:, 5, :, :])

        outputs_resized = F.interpolate(outputs, size=(TARGET_HEIGHT, TARGET_WIDTH), mode='bilinear', align_corners=False)
        labels = to_one_hot(labels, n_class)
        
        loss = criterion(outputs_resized, labels.long())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()

    test_loss = 0.0
    dimssnet.eval()
    with torch.no_grad():
        for images, labels, _ in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            
            images = [image.unsqueeze(1) for image in images]
            images = torch.stack(images, dim=0)
            images = images.float()
            
            outputs = dimssnet(images[:, 0, :, :], images[:, 1, :, :], images[:, 2, :, :],
                                images[:, 3, :, :], images[:, 4, :, :], images[:, 5, :, :])
            
            outputs_resized = F.interpolate(outputs, size=(TARGET_HEIGHT, TARGET_WIDTH), mode='bilinear', align_corners=False)
            labels = to_one_hot(labels, n_class)
            
            loss = criterion(outputs_resized, labels.long())
            test_loss += loss.item()
    
    test_loss_record.append(test_loss/len(test_loader))
    train_loss_record.append(running_loss/len(train_loader))
    print(f"Epoch {epoch+1}/{num_epochs}, Training Loss: {running_loss/len(train_loader):.4f}, Test Loss: {test_loss/len(test_loader):.4f}")

# Create a new state_dict and add the "model" prefix, in order to match the loading method in metrics
new_state_dict = {}
for k, v in dimssnet.state_dict().items():
    new_key = f"model.{k}" 
    new_state_dict[new_key] = v

print("dimssnet successfully saved!")    
print("Training finished.")
end_time = time.time()
total_time = end_time - start_time
print(f"Training completed in {total_time} seconds")

Inference

# Run inference on validation set and testing set
import pandas as pd

def run_inference(model, data_loader):
    model.eval()
    predictions = []
    filenames = []
    
    with torch.no_grad():
        for images, file_names in data_loader:
            images = images.cuda() if torch.cuda.is_available() else images
            
            x = [image.unsqueeze(1) for image in images]
            x = torch.stack(x, dim=0)
            x = x.float()
            
            # Forward propagation
            outputs = model(x[:, 0, :, :], x[:, 1, :, :], x[:, 2, :, :],
                            x[:, 3, :, :], x[:, 4, :, :], x[:, 5, :, :])
            
            # Adjust the output size
            outputs = F.interpolate(outputs, size=(50, 181), mode='bilinear', align_corners=False)

            preds = torch.argmax(outputs, dim=1)
            
            preds = preds - 1
            
            for i, pred in enumerate(preds):
                predictions.append(pred.cpu().numpy().flatten())
                filenames.append(file_names[i])
    
    return predictions, filenames

# Load validation set
if os.environ.get('DATA_PATH'):
    DATA_PATH = os.environ.get("DATA_PATH") + "/" 
else:
    DATA_PATH = ""  # Fallback for local testing
# Load validation set
val_paths = generate_file_paths(DATA_PATH + 'validation_set')
val_dataset = CustomDataset_test(file_paths=val_paths)
val_loader = DataLoader(
    val_dataset,
    batch_size=1,
    shuffle=False,
    num_workers=0
)

# Load testing set
if os.environ.get('DATA_PATH'):
    DATA_PATH = os.environ.get("DATA_PATH") + "/" 
else:
    DATA_PATH = ""  # Fallback for local testing
# Load testing set
test_paths = generate_file_paths(DATA_PATH + 'testing_set')
test_dataset = CustomDataset_test(file_paths=test_paths)
test_loader = DataLoader(
    test_dataset,
    batch_size=1,
    shuffle=False,
    num_workers=0
)

# Run inference on validation set
print("Running inference on validation set...")
val_predictions, val_filenames = run_inference(dimssnet, val_loader)

# Save validation results to CSV
val_results = []
for filename, pred in zip(val_filenames, val_predictions):
    # Create a row with filename and flattened predictions
    row = {'filename': filename}
    for i, p in enumerate(pred):
        row[f'pixel_{i}'] = p
    val_results.append(row)

val_df = pd.DataFrame(val_results)
val_df.to_csv('submission_val.csv', index=False)
print(f"Validation results saved to output_validation.csv with shape: {val_df.shape}")

# Run inference on testing set
print("Running inference on testing set...")
test_predictions, test_filenames = run_inference(dimssnet, test_loader)

# Save testing results to CSV
test_results = []
for filename, pred in zip(test_filenames, test_predictions):
    # Create a row with filename and flattened predictions
    row = {'filename': filename}
    for i, p in enumerate(pred):
        row[f'pixel_{i}'] = p
    test_results.append(row)

test_df = pd.DataFrame(test_results)
test_df.to_csv('submission_test.csv', index=False)
print(f"Testing results saved to output_testing.csv with shape: {test_df.shape}")

print("\nInference completed! Results saved to:")
print("- submission_val_ref.csv (for validation set leaderboard)")
print("- submission_test_ref.csv (for testing set leaderboard)")

import zipfile
import os

# Define the files to zip and the zip file name.
files_to_zip = ['submission_val.csv', 'submission_test.csv']
zip_filename = 'submission.zip'

# Create a zip file
with zipfile.ZipFile(zip_filename, 'w') as zipf:
    for file in files_to_zip:
        # Add the file to the zip fil
        zipf.write(file, os.path.basename(file))

print(f'{zip_filename} Created successfully!')