def test(frame, opt, rgb, transform_test, num_workers, label_indice,
model_path):
img = Image.fromarray(frame)
testset = Countmap_Dataset(img,rgb,transform=transform_test,\
if_test=True, IF_loadmem=opt['IF_savemem_test'])
testloader = DataLoader(testset,
batch_size=opt['test_batch_size'],
shuffle=False,
num_workers=num_workers)
# init networks
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice) + 1
div_times = opt['div_times']
net = SSDCNet_classify(class_num,label_indice,div_times=div_times,\
frontend_name='VGG16',block_num=5,\
IF_pre_bn=False,IF_freeze_bn=False,load_weights=True,\
psize=opt['psize'],pstride = opt['pstride'],parse_method ='maxp').cuda()
# test the min epoch
mod_path = 'best_epoch.pth'
mod_path = os.path.join(opt['model_path'], mod_path)
if os.path.exists(mod_path):
all_state_dict = torch.load(mod_path)
net.load_state_dict(all_state_dict['net_state_dict'])
tmp_epoch_num = all_state_dict['tmp_epoch_num']
log_save_path = os.path.join(
model_path, 'log-epoch-min[%d]-%s.txt' %
(tmp_epoch_num + 1, opt['parse_method']))
# test
test_log, count = test_phase(opt,
net,
testloader,
log_save_path=log_save_path)
return count
def main(opt):
# save folder
save_folder = opt['model_path']
# =============================================================================
# inital setting
# =============================================================================
# 1.Initial setting
# --1.1 dataset setting
dataset = opt['dataset']
root_dir = opt['root_dir']
num_workers = opt['num_workers']
img_subsubdir = 'images'
tar_subsubdir = 'gtdens'
transform_test = []
# --1.3 use initial setting to generate
# set label_indice
label_indice = np.arange(opt['step'], opt['max_num'] + opt['step'],
opt['step'])
add = np.array(
[1e-6, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45])
label_indice = np.concatenate((add, label_indice))
print(label_indice)
opt['label_indice'] = label_indice
opt['class_num'] = label_indice.size + 1
#test settings
img_dir = os.path.join(root_dir, 'test', img_subsubdir)
tar_dir = os.path.join(root_dir, 'test', tar_subsubdir)
rgb_dir = os.path.join(root_dir, 'rgbstate.mat')
testset = Countmap_Dataset(img_dir,tar_dir,rgb_dir,transform=transform_test,\
if_test=True, IF_loadmem=opt['IF_savemem_test'])
testloader = DataLoader(testset,
batch_size=opt['test_batch_size'],
shuffle=False,
num_workers=num_workers)
# init networks
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice) + 1
div_times = opt['div_times']
net = SSDCNet_classify(class_num,label_indice,div_times=div_times,\
frontend_name='VGG16',block_num=5,\
IF_pre_bn=False,IF_freeze_bn=False,load_weights=True,\
psize=opt['psize'],pstride = opt['pstride'],parse_method ='maxp').cuda()
# test the min epoch
mod_path = 'best_epoch.pth'
mod_path = os.path.join(opt['model_path'], mod_path)
if os.path.exists(mod_path):
all_state_dict = torch.load(mod_path)
net.load_state_dict(all_state_dict['net_state_dict'])
tmp_epoch_num = all_state_dict['tmp_epoch_num']
log_save_path = os.path.join(save_folder,'log-epoch-min[%d]-%s.txt' \
%(tmp_epoch_num+1,opt['parse_method']) )
# test
test_log = test_phase(opt,
net,
testloader,
log_save_path=log_save_path)
def __init__(self, medium):
# Using OpenCV to capture from device 0. If you have trouble capturing
# from a webcam, comment the line below out and use a video file
# instead.
if medium == "video":
self.stream = FileVideoStream("videos/Assult_attack.mp4").start()
elif medium == "webcam":
self.stream = WebcamVideoStream(src=0).start()
save_folder = 'model'
root_dir = r'data'
num_workers = 1
transform_test = []
label_indice = np.arange(0.5, 7.5, 0.5)
add = np.array(
[1e-6, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45])
label_indice = np.concatenate((add, label_indice))
print("label_indice:", label_indice)
rgb_dir = os.path.join(root_dir, 'rgbstate.mat')
mat = sio.loadmat(rgb_dir)
self.rgb = mat['rgbMean'].reshape(1, 1, 3)
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice) + 1
self.net = SSDCNet_classify(class_num,
label_indice,
div_times=2,
frontend_name='VGG16',
block_num=5,
IF_pre_bn=False,
IF_freeze_bn=False,
load_weights=True,
psize=64,
pstride=64,
parse_method='maxp').cuda()
mod_path = 'best_epoch.pth'
mod_path = os.path.join('model', mod_path)
print(mod_path)
if os.path.exists(mod_path):
all_state_dict = torch.load(mod_path)
print(mod_path)
self.net.load_state_dict(all_state_dict['net_state_dict'])
tmp_epoch_num = all_state_dict['tmp_epoch_num']
log_save_path = os.path.join(
save_folder,
'log-epoch-min[%d]-%s.txt' % (tmp_epoch_num + 1, 'maxp'))
print("end of function")
print("Please enter file name", file=sys.stderr)
exit()
mod_path = 'best_epoch.pth'
max_num = 7
step = 0.5
label_indice = np.arange(step, max_num + step, step)
add = np.array(
[1e-6, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45])
label_indice = np.concatenate((add, label_indice))
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice) + 1
div_times = 2
psize, pstride = 64, 64
if cuda:
net = SSDCNet_classify(class_num,label_indice,div_times=div_times,\
frontend_name='VGG16',block_num=5,\
IF_pre_bn=False,IF_freeze_bn=False,load_weights=True,\
psize=psize,pstride = pstride,parse_method ='maxp').cuda()
else:
net = SSDCNet_classify(class_num,label_indice,div_times=div_times,\
frontend_name='VGG16',block_num=5,\
IF_pre_bn=False,IF_freeze_bn=False,load_weights=True,\
psize=psize,pstride = pstride,parse_method ='maxp').cpu()
if not os.path.isabs(mod_path):
mod_path = os.path.join(FILE_DIR, mod_path)
if os.path.exists(mod_path):
if cuda:
all_state_dict = torch.load(mod_path)
else:
all_state_dict = torch.load(mod_path, map_location='cpu')
net.load_state_dict(all_state_dict['net_state_dict'])
tmp_epoch_num = all_state_dict['tmp_epoch_num']
def main(model_path, input_source, SDCNet=False):
if os.path.isfile(input_source):
video_name = input_source.split("/")[-1]
video_path = input_source.replace(video_name,"")
print("Processing video {} in {}".format(colored(video_name, 'red'), colored(video_path, 'green')))
cap = cv2.VideoCapture(input_source)
else:
cap = VideoStream(input_source).start()
# --1.2 use initial setting to generate
# set label_indice
label_indice = np.arange(0.5, 22+0.5/2, 0.5)
add = np.array([1e-6, 0.05, 0.10, 0.15, 0.20,
0.25, 0.30, 0.35, 0.40, 0.45])
label_indice = np.concatenate((add, label_indice))
# init networks
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice)+1
div_times = 2
if SDCNet:
net = SDCNet_VGG16_classify(
class_num, label_indice, load_weights=True)
else:
net = SSDCNet_classify(class_num, label_indice, div_times=div_times,
frontend_name='VGG16', block_num=5,
IF_pre_bn=False, IF_freeze_bn=False, load_weights=True,
parse_method='maxp')
if os.path.exists(model_path):
print("Adding Weights ....")
all_state_dict = torch.load(model_path,map_location=torch.device('cpu'))
net.load_state_dict(all_state_dict['net_state_dict'])
net.eval()
else:
print("Can't find trained weights!!")
exit()
first_frame = True
while True:
flag, image = cap.read()
output_image = np.copy(image)
if first_frame:
rois = cv2.selectROIs("frame", image, False, False)
first_frame = False
# print(roi)
if flag:
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
else:
break
sum = 0
for roi in rois:
roi_image = image[int(roi[1]):int(roi[1]+roi[3]),
int(roi[0]):int(roi[0]+roi[2])]
roi_image = cv2.resize(roi_image, (256, 256))
roi_image = np.transpose(roi_image, (2, 0, 1))
roi_image = torch.Tensor(roi_image[None, :, :, :])
# w = image.shape[-1]
# h = image.shape[-2]
# pad_w = 64 - w%64
# padding_left = int(pad_w/2)
# padding_right = pad_w - padding_left
# pad_h = 64 - h%64
# padding_top = int(pad_h/2)
# padding_bottom = pad_h - padding_top
# image = torch.nn.functional.pad(image, (padding_left, padding_right, padding_top, padding_bottom))
with torch.no_grad():
features = net(roi_image)
div_res = net.resample(features)
merge_res = net.parse_merge(div_res)
if SDCNet:
outputs = merge_res['div'+str(net.args['div_times'])]
else:
outputs = merge_res['div'+str(net.div_times)]
del merge_res
cv2.rectangle(output_image, (int(roi[0]), int(roi[1])), (int(
roi[0]+roi[2]), int(roi[1]+roi[3])), (255, 0, 0), thickness=3)
sum += int(outputs.sum().item())
cv2.putText(output_image, "{}".format(sum),
(30, 50), cv2.FONT_HERSHEY_PLAIN, 2,
(255, 0, 0), 3)
cv2.imshow("frame", output_image)
if cv2.waitKey(1) == ord('q'):
cap.release()
exit()
cap.release()
class VideoCamera(object):
def __init__(self, medium):
# Using OpenCV to capture from device 0. If you have trouble capturing
# from a webcam, comment the line below out and use a video file
# instead.
if medium == "video":
self.stream = FileVideoStream("videos/Assult_attack.mp4").start()
elif medium == "webcam":
self.stream = WebcamVideoStream(src=0).start()
save_folder = 'model'
root_dir = r'data'
num_workers = 1
transform_test = []
label_indice = np.arange(0.5, 7.5, 0.5)
add = np.array(
[1e-6, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45])
label_indice = np.concatenate((add, label_indice))
print("label_indice:", label_indice)
rgb_dir = os.path.join(root_dir, 'rgbstate.mat')
mat = sio.loadmat(rgb_dir)
self.rgb = mat['rgbMean'].reshape(1, 1, 3)
label_indice = torch.Tensor(label_indice)
class_num = len(label_indice) + 1
self.net = SSDCNet_classify(class_num,
label_indice,
div_times=2,
frontend_name='VGG16',
block_num=5,
IF_pre_bn=False,
IF_freeze_bn=False,
load_weights=True,
psize=64,
pstride=64,
parse_method='maxp').cuda()
mod_path = 'best_epoch.pth'
mod_path = os.path.join('model', mod_path)
print(mod_path)
if os.path.exists(mod_path):
all_state_dict = torch.load(mod_path)
print(mod_path)
self.net.load_state_dict(all_state_dict['net_state_dict'])
tmp_epoch_num = all_state_dict['tmp_epoch_num']
log_save_path = os.path.join(
save_folder,
'log-epoch-min[%d]-%s.txt' % (tmp_epoch_num + 1, 'maxp'))
print("end of function")
def __del__(self):
self.stream.stop()
def get_pad(self, inputs, DIV=64):
h, w = inputs.size()[-2:]
ph, pw = (DIV - h % DIV), (DIV - w % DIV)
# print(ph,pw)
tmp_pad = [0, 0, 0, 0]
if (ph != DIV):
tmp_pad[2], tmp_pad[3] = 0, ph
if (pw != DIV):
tmp_pad[0], tmp_pad[1] = 0, pw
# print(tmp_pad)
inputs = F.pad(inputs, tmp_pad)
return inputs
def get_frame(self):
image = self.stream.read()
# We are using Motion JPEG, but OpenCV defaults to capture raw images,
# so we must encode it into JPEG in order to correctly display the
# video stream.
print(image.shape)
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img).convert('RGB')
print(type(img))
img = img.resize((320, 240))
#img.show()
img = transforms.ToTensor()(img)
img = self.get_pad(img, DIV=64)
img = img - torch.Tensor(self.rgb).view(3, 1, 1)
# input_file.save(input_file.filename)
# print(input_file.filename)
inputs = img.type(torch.float32)
inputs = inputs.unsqueeze(0)
inputs = inputs.cuda()
print(inputs.size())
features = self.net(inputs)
div_res = self.net.resample(features)
merge_res = self.net.parse_merge(div_res)
outputs = merge_res['div' + str(self.net.div_times)]
del merge_res
pre = (outputs).sum()
threshold = 20
result = int(math.floor(pre.item()))
if result <= threshold:
result = str(int(math.floor(pre.item()))) + "person "
else:
result = str(int(math.floor(pre.item()))) + \
"person. " + "Please distance!"
print('pre:', pre)
startX = int(0)
startY = int(0)
endX = int(75)
endY = int(75)
cv2.putText(image, result, (endX, endY), cv2.FONT_HERSHEY_SIMPLEX, 2,
(0, 255, 0), 2)
ret, jpeg = cv2.imencode('.jpg', image)
data = []
data.append(jpeg.tobytes())
data.append(result)
return data