forked from koyeongmin/PINet_new
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hourglass_network.py
37 lines (31 loc) · 1.21 KB
/
hourglass_network.py
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
#########################################################################
##
## Structure of network.
##
#########################################################################
import torch
import torch.nn as nn
from util_hourglass import *
####################################################################
##
## lane_detection_network
##
####################################################################
class lane_detection_network(nn.Module):
def __init__(self):
super(lane_detection_network, self).__init__()
self.resizing = resize_layer(3, 128)
#feature extraction
self.layer1 = hourglass_block(128, 128)
self.layer2 = hourglass_block(128, 128)
self.layer3 = hourglass_block(128, 128)
self.layer4 = hourglass_block(128, 128)
def forward(self, inputs):
#feature extraction
out = self.resizing(inputs)
result1, out, feature1 = self.layer1(out)
result2, out, feature2 = self.layer2(out)
result3, out, feature3 = self.layer3(out)
result4, out, feature4 = self.layer4(out)
return [result1, result2, result3, result4], [feature1, feature2, feature3, feature4]
#return [result1], [feature1]