Views
No views yet
model.pth — PyTorch checkpointmodel.safetensors — Safetensors weightsmodel.onnx — ONNX (opset 17, dynamic axes)model_fp16.plan — TensorRT FP16model_fp32.plan — TensorRT FP321import torch
2from safetensors.torch import load_file
3from def_dhif.models.unet_dhif import UNetDHiF
4
5model = UNetDHiF(in_channels=1, out_channels=1, base_ch=64)
6state = load_file("model.safetensors")
7model.load_state_dict(state)
8model.eval()
9
10# Single-channel IR image, padded to 512x512
11x = torch.randn(1, 1, 512, 512)
12pred = model(x) # Binary segmentation mask