class Detector(nn.Module):
def __init__(self, config_path, weights_path, input_size=None, conf_thresh=0.5, nms_thresh=0.4):
super(Detector, self).__init__()
self.input_size = input_size
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
# Initialize Darknet for detection
self.model = Darknet(config_path, input_size=input_size)
self.model.load_weights(weights_path)
self.model.eval()
def device(self):
return next(self.model.parameters()).device
def forward(self, frame, swapRB=False):
x = image_to_tensor(frame, swapRB)
_, _, fh, fw = x.size()
device = self.device()
x = x.to(device)
x = letterbox_resize(x, self.input_size, constant_value=127.5)
x = x / 255.0
with torch.no_grad():
y = self.model.forward(x)
output = []
for i, prediction in enumerate(y): # Enumerate on batch
detection = non_max_suppression(prediction.cpu(), self.conf_thresh, self.nms_thresh)
if detection is not None:
detection = bbox_fit(detection, (fh, fw), self.input_size).to(device)
output.append(detection)
return output
def update(self, conf_thresh=None, nms_thresh=None, weights_path=None):
if conf_thresh is not None:
self.conf_thresh = conf_thresh
if nms_thresh is not None:
self.nms_thresh = nms_thresh
if weights_path is not None:
device = self.device()
self.model.cpu().load_weights(weights_path)
self.model.to(device)
CUDA = torch.cuda.is_available()
num_classes = 80
bbox_attrs = 5 + num_classes
print("Loading network.....")
model = Darknet(args.cfgfile)
model.load_weights(args.weightsfile)
print("Network successfully loaded")
model.net_info["height"] = args.reso
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
if CUDA:
model.cpu().half()
model(get_test_input(inp_dim, CUDA), CUDA)
model.eval()
videofile = 'video.avi'
cap = cv2.VideoCapture(videofile)
assert cap.isOpened(), 'Cannot capture source'
frames = 0
start = time.time()
while cap.isOpened():
classes = load_classes('data/coco.names')
#Set up the neural network
print("Loading network.....")
model = Darknet("cfg/yolov3-spp.cfg")
model.load_weights("yolov3-spp.weights")
print("Network successfully loaded")
model.net_info["height"] = "608"
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
#If there's a GPU availible, put the model on GPU
if CUDA:
model.cpu()
#Set the model in evaluation mode
model.eval()
#Detection phase
try:
imlist = []
imlist.append(osp.join(osp.realpath('.'), images))
except FileNotFoundError:
print("No file or directory with the name {}".format(images))
exit()
batches = list(
map(prep_image, imlist, [inp_dim for x in range(len(imlist))]))
im_batches = [x[0] for x in batches]
kwargs = {'num_workers': num_workers, 'pin_memory': True} if use_cuda else {}
# Get the iterable test dataset containing input and targets
test_loader = torch.utils.data.DataLoader(
dataset.listDataset(testlist, shape=(init_width, init_height),
shuffle=False,
transform=transforms.Compose([
transforms.ToTensor(), # These transformations are executed to every sample
]), train=True),
batch_size=batch_size, shuffle=False, **kwargs)
if use_cuda:
if ngpus > 1:
model = torch.nn.DataParallel(model).cuda()
else:
# model = model.cuda() # Original
model = model.cpu()
params_dict = dict(model.named_parameters()) # All the parameters in the model
params = []
# Changing the hyper-parameters of parameters
for key, value in params_dict.items():
if key.find('.bn') >= 0 or key.find('.bias') >= 0:
params += [{'params': [value], 'weight_decay': 0.0}]
else:
params += [{'params': [value], 'weight_decay': decay*batch_size}]
# Stochastic Gradient Descent
# optimizer = optim.SGD(model.parameters(), lr=learning_rate/batch_size, momentum=momentum, dampening=0, weight_decay=decay*batch_size)
# RMSProp
optimizer = optim.RMSprop(model.parameters(), lr=learning_rate/batch_size, momentum=momentum, weight_decay=decay*batch_size)
