Spaces:
Build error
Build error
File size: 2,711 Bytes
4d85df4 |
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 |
# encoding: utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from monoscene.modules import SegmentationHead
from monoscene.CRP3D import CPMegaVoxels
from monoscene.modules import Process, Upsample, Downsample
class UNet3D(nn.Module):
def __init__(
self,
class_num,
norm_layer,
full_scene_size,
feature,
project_scale,
context_prior=None,
bn_momentum=0.1,
):
super(UNet3D, self).__init__()
self.business_layer = []
self.project_scale = project_scale
self.full_scene_size = full_scene_size
self.feature = feature
size_l1 = (
int(self.full_scene_size[0] / project_scale),
int(self.full_scene_size[1] / project_scale),
int(self.full_scene_size[2] / project_scale),
)
size_l2 = (size_l1[0] // 2, size_l1[1] // 2, size_l1[2] // 2)
size_l3 = (size_l2[0] // 2, size_l2[1] // 2, size_l2[2] // 2)
dilations = [1, 2, 3]
self.process_l1 = nn.Sequential(
Process(self.feature, norm_layer, bn_momentum, dilations=[1, 2, 3]),
Downsample(self.feature, norm_layer, bn_momentum),
)
self.process_l2 = nn.Sequential(
Process(self.feature * 2, norm_layer, bn_momentum, dilations=[1, 2, 3]),
Downsample(self.feature * 2, norm_layer, bn_momentum),
)
self.up_13_l2 = Upsample(
self.feature * 4, self.feature * 2, norm_layer, bn_momentum
)
self.up_12_l1 = Upsample(
self.feature * 2, self.feature, norm_layer, bn_momentum
)
self.up_l1_lfull = Upsample(
self.feature, self.feature // 2, norm_layer, bn_momentum
)
self.ssc_head = SegmentationHead(
self.feature // 2, self.feature // 2, class_num, dilations
)
self.context_prior = context_prior
if context_prior:
self.CP_mega_voxels = CPMegaVoxels(
self.feature * 4, size_l3, bn_momentum=bn_momentum
)
def forward(self, input_dict):
res = {}
x3d_l1 = input_dict["x3d"]
x3d_l2 = self.process_l1(x3d_l1)
x3d_l3 = self.process_l2(x3d_l2)
if self.context_prior:
ret = self.CP_mega_voxels(x3d_l3)
x3d_l3 = ret["x"]
for k in ret.keys():
res[k] = ret[k]
x3d_up_l2 = self.up_13_l2(x3d_l3) + x3d_l2
x3d_up_l1 = self.up_12_l1(x3d_up_l2) + x3d_l1
x3d_up_lfull = self.up_l1_lfull(x3d_up_l1)
ssc_logit_full = self.ssc_head(x3d_up_lfull)
res["ssc_logit"] = ssc_logit_full
return res
|