def forward(self, x, inp_dim, num_classes, confidence):
x = x.data
global CUDA
prediction = x
prediction = predict_transform(prediction, inp_dim, self.anchors,
num_classes, confidence, CUDA)
return prediction
def forward(self, x, CUDA):
detections = []
modules = self.blocks[1:]
outputs = {} #We cache the outputs for the route layer
write = 0
for i in range(len(modules)):
module_type = (modules[i]["type"])
if module_type == "convolutional" or module_type == "upsample" or module_type == "maxpool":
x = self.module_list[i](x)
outputs[i] = x
elif module_type == "route":
layers = modules[i]["layers"]
layers = [int(a) for a in layers]
if (layers[0]) > 0:
layers[0] = layers[0] - i
if len(layers) == 1:
x = outputs[i + (layers[0])]
else:
if (layers[1]) > 0:
layers[1] = layers[1] - i
map1 = outputs[i + layers[0]]
map2 = outputs[i + layers[1]]
x = torch.cat((map1, map2), 1)
outputs[i] = x
elif module_type == "shortcut":
from_ = int(modules[i]["from"])
x = outputs[i - 1] + outputs[i + from_]
outputs[i] = x
elif module_type == 'yolo':
anchors = self.module_list[i][0].anchors
#Get the input dimensions
inp_dim = int(self.net_info["height"])
#Get the number of classes
num_classes = int(modules[i]["classes"])
#Output the result
x = x.data
x = predict_transform(x, inp_dim, anchors, num_classes, CUDA)
if type(x) == int:
continue
if not write:
detections = x
write = 1
else:
detections = torch.cat((detections, x), 1)
outputs[i] = outputs[i - 1]
try:
return detections
except:
return 0
