def detect_cv2_camera(cfgfile, weightfile):
import cv2
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture("./test.mp4")
cap.set(3, 1280)
cap.set(4, 720)
print("Starting the YOLO loop...")
num_classes = m.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/x.names'
class_names = load_class_names(namesfile)
while True:
ret, img = cap.read()
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
finish = time.time()
print('Predicted in %f seconds.' % (finish - start))
result_img = plot_boxes_cv2(img,
boxes[0],
savename=None,
class_names=class_names)
cv2.imshow('Yolo demo', result_img)
cv2.waitKey(1)
cap.release()
def detect_cv2_img(cfgfile, weightfile, img_file):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
# print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
img = cv2.imread(img_file)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
boxes_darknet_format = to_darknet_format(boxes, img.shape[0], img.shape[1])
print(boxes[0])
print(boxes_darknet_format)
return boxes[0]
def transform_to_onnx(cfgfile, weightfile, batch_size=1):
model = Darknet(cfgfile)
model.print_network()
model.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
dynamic = False
if batch_size <= 0:
dynamic = True
input_names = ["input"]
output_names = ['boxes', 'confs']
if dynamic:
##############################################################################
x = torch.randn((1, model.height, model.width, 3), requires_grad=True).byte()
##############################################################################
#x = torch.randn((1, 3, model.height, model.width), requires_grad=True)
onnx_file_name = "yolov4_-1_3_{}_{}_dynamic.onnx".format(model.height, model.width)
dynamic_axes = {"input": {0: "batch_size"}, "boxes": {0: "batch_size"}, "confs": {0: "batch_size"}}
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=12, #11
do_constant_folding=True,
input_names=input_names, output_names=output_names,
dynamic_axes=dynamic_axes)
print('Onnx model exporting done')
return onnx_file_name
else:
##############################################################################
x = torch.randn((batch_size, model.height, model.width, 3), requires_grad=True).byte()
##############################################################################
#x = torch.randn((batch_size, 3, model.height, model.width), requires_grad=True)
onnx_file_name = "yolov4_{}_3_{}_{}_static.onnx".format(batch_size, model.height, model.width)
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=input_names, output_names=output_names,
dynamic_axes=None)
print('Onnx model exporting done')
return onnx_file_name
def detect_imges(cfgfile,
weightfile,
imgfile_list=[
os.path.dirname(__file__) + '/data/dog.jpg',
os.path.dirname(__file__) + '/data/giraffe.jpg'
]):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
num_classes = 80
if num_classes == 20:
namesfile = os.path.dirname(__file__) + '/data/voc.names'
elif num_classes == 80:
namesfile = os.path.dirname(__file__) + '/data/coco.names'
else:
namesfile = os.path.dirname(__file__) + '/data/names'
use_cuda = 0
if use_cuda:
m.cuda()
imges = []
imges_list = []
for imgfile in imgfile_list:
img = Image.open(imgfile).convert('RGB')
imges_list.append(img)
sized = img.resize((m.width, m.height))
imges.append(np.expand_dims(np.array(sized), axis=0))
images = np.concatenate(imges, 0)
for i in range(2):
start = time.time()
boxes = do_detect(m, images, 0.5, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
for i, (img, box) in enumerate(zip(imges_list, boxes)):
plot_boxes(img, box, 'predictions{}.jpg'.format(i), class_names)
def detect_cv2(cfgfile, weightfile, img):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
num_classes = m.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/x.names'
class_names = load_class_names(namesfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
objects = []
for i, box in enumerate(boxes[0]):
dic = {}
dic['kind'] = class_names[box[6]]
dic['confidence'] = box[4]
dic.update(get_bbox_coordinates(img, box))
cropped_img = crop_box(img, box)
dic['feature'] = extract_feature(cropped_img)
objects.append(dic)
finish = time.time()
print('Predicted in %f seconds.' % (finish - start))
plot_boxes_cv2(img, boxes[0], savename='predictions.jpg', class_names=class_names)
return({'objects': objects})
def detect_cv2(cfgfile, weightfile, imgfile):
import cv2
m = Darknet(cfgfile) # 创建 Darknet 模型对象 m
m.print_network() # 打印网络结构信息
m.load_weights(weightfile) # 加载网络权重值 在 tools/darknet2pytorch.py 函数中
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda() # 如果使用 cuda，则将模型对象拷贝至显存
num_classes = m.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/x.names'
class_names = load_class_names(namesfile) # 加载类别名
# 如果用 PIL 打开图像
# img = Image.open(imgfile).convert('RGB')
# sized = img.resize((m.width, m.height))
img = cv2.imread(imgfile)
cv2.imwrite('./debug/img.jpg', img)
# print(m.width, m.height) # (608, 608)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6,
use_cuda) # 做检测，返回的 boxes 是最晚 NMS 后的检测框
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img,
boxes[0],
savename='./debug/predictions.jpg',
class_names=class_names) # raw
def detect(cfgfile, weightfile, imgfile):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
img = Image.open(imgfile).convert('RGB')
sized = img.resize((m.width, m.height))
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.5, num_classes, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
plot_boxes(img, boxes, 'predictions.jpg', class_names)
def detect(cfgfile, weightfile, imgfile):
m = Darknet(cfgfile)
checkpoint = torch.load(weightfile)
model_dict = m.state_dict()
pretrained_dict = checkpoint
keys = []
for k, v in pretrained_dict.items():
keys.append(k)
i = 0
for k, v in model_dict.items():
if v.size() == pretrained_dict[keys[i]].size():
model_dict[k] = pretrained_dict[keys[i]]
i = i + 1
m.load_state_dict(model_dict)
# m.load_state_dict(torch.load(weightfile))
# m.print_network()
# m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
namesfile = 'data/mydata.names'
use_cuda = 1
if use_cuda:
m.cuda()
input_img = cv2.imread(imgfile)
# orig_img = Image.open(imgfile).convert('RGB')
start = time.time()
boxes, scale = do_detect(m, input_img, 0.5, 0.4, use_cuda)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
# draw_boxes(input_img,boxes,scale=scale)
plot_boxes_cv2(input_img,
boxes,
'predictions1.jpg',
class_names=class_names,
scale=scale)
def detect_BEV_flat(cfgfile, weightfile, imgfile):
"""Detect elements in BEV map with yolov4_BEV_flat
Args:
cfgfile (str): Path to .cfg file
weightfile (str): Path to .weights file
imgfile (str): Path to image on which we want to run BEV detection
"""
# load model
m = Darknet(cfgfile, model_type="BEV_flat")
m.print_network()
m.load_weights(weightfile, cut_off=54)
print("Loading backbone from %s... Done!" % (weightfile))
# push to GPU
if use_cuda:
m.cuda()
# load names
namesfile = "names/BEV.names"
class_names = load_class_names(namesfile)
# read sample image
img = cv2.imread(imgfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
# create batch
sized = np.expand_dims(sized, 0)
sized = np.concatenate((sized, sized), 0)
# run inference
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
finish = time.time()
print("%s: Predicted in %f seconds." % (imgfile, (finish - start)))
def detect_cv2(cfgfile, weightfile, imgfile):
import cv2
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
num_classes = m.num_classes
num_classes = 1
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/head.names'
class_names = load_class_names(namesfile)
img = cv2.imread(imgfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
print(boxes)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img,
boxes[0],
savename='predictions.jpg',
class_names=class_names)
def detect(json_dir, video_dir, save_dir):
starttime = timeit.default_timer()
Path(save_dir).mkdir(parents=True, exist_ok=True)
cfgfile = config['detector']['cfgfile']
weightfile = config['detector']['weightfile']
model = Darknet(cfgfile)
model.load_weights(weightfile)
model.cuda()
class_names = config['detector']['originclassnames']
cam_datas = get_list_data(json_dir)
for cam_data in cam_datas:
cam_name = cam_data['camName']
roi_poly = Polygon(cam_data['shapes'][0]['points'])
video_path = os.path.join(video_dir, cam_name + '.mp4')
video_cap = cv2.VideoCapture(video_path)
num_frames = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
imgs = []
for i in tqdm(range(num_frames),
desc='Extracting {}'.format(cam_name)):
success, img = video_cap.read()
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
imgs.append(img)
boxes = detect_yolo(model, class_names, imgs, cam_name,
config['detector']['batchsize'])
# remove bboxes out of MOI
if config['remove_not_intersec_moi']:
boxes = [
check_intersect_box(box_list, roi_poly) for box_list in boxes
]
if save_dir:
filepath = os.path.join(save_dir, cam_name)
boxes = np.array(boxes)
np.save(filepath, boxes)
endtime = timeit.default_timer()
print('Detect time: {} seconds'.format(endtime - starttime))
def detect_cv2(cfgfile, weightfile, imgfile, outfile):
import cv2
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
num_classes = m.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'D:/work_source/CV_Project/datasets/xi_an_20201125/all/names_xi_an_20201125.txt'
class_names = load_class_names(namesfile)
print('imgfile: ', imgfile)
img = cv2.imread(imgfile)
print(m.width, m.height)
# print('demo pic size:', img.shape)
sized = cv2.resize(img, (m.width, m.height))
# print('demo pic resize to:',sized.shape)
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
boxes = []
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img, boxes[0], savename=outfile, class_names=class_names)
def detect_skimage(cfgfile, weightfile, imgfile):
from skimage import io
from skimage.transform import resize
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
num_classes = m.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/x.names'
class_names = load_class_names(namesfile)
img = io.imread(imgfile)
sized = resize(img, (m.width, m.height)) * 255
boxes = []
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img,
boxes,
savename='predictions3.jpg',
class_names=class_names)
def fransform_to_onnx(cfgfile, weightfile, batch_size=1):
model = Darknet(cfgfile)
model.print_network()
model.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
# model.cuda()
x = torch.randn((batch_size, 3, model.height, model.width), requires_grad=True) # .cuda()
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
"yolov4_{}_3_{}_{}.onnx".format(batch_size, model.height, model.width),
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=['input'], output_names=['output_1', 'output_2', 'output_3'],
dynamic_axes=None)
print('Onnx model exporting done')
def detect_cv2(cfgfile, weightfile, imgfiles, namesfile):
CENTER_X = 0
CENTER_Y = 1
WIDTH = 2
HEIGHT = 3
CONFIDENCE = 4
CLASS_ID = 6
m = Darknet(cfgfile)
# m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
num_classes = m.num_classes
class_names = load_class_names(namesfile)
f = open(imgfiles, 'r')
filelist = f.read().splitlines()
f.close()
# print(filelist)
start_time = time.time()
submit_results = []
pbar = tqdm(total=len(filelist))
for idx, imgfile in enumerate(filelist):
# print(imgfile)
img = cv2.imread(imgfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
img_height, img_width, _ = img.shape
predict = do_detect(m, sized, 0.2, 0.6, use_cuda)[0]
pbar.update(1)
result = {}
result['bbox'] = [
output_transform(img_width, img_height, box[CENTER_X],
box[CENTER_Y], box[WIDTH], box[HEIGHT])
for box in predict
]
# print(result['bbox'])
result['score'] = [float(box[CONFIDENCE]) for box in predict]
# print(result['score'])
result['label'] = [
int(box[CLASS_ID]) if not box[CLASS_ID] == 0 else 10
for box in predict
]
# print(result['label'])
# print(result)
submit_results.append(result)
end_time = time.time()
print('number of predicts: {}'.format(len(submit_results)))
print('fps: {}'.format(len(submit_results) / (end_time - start_time)))
print('speed: {} ms per image'.format(
(end_time - start_time) * 1000 / len(submit_results)))
print('output results for submission: {}'.format('submission.json'))
with open('submission.json', 'w') as outfile:
json.dump(submit_results, outfile)
dest='cfgfile')
parser.add_argument('-weightfile',
type=str,
default='./checkpoints/Yolov4_epoch1.pth',
help='path of trained model.',
dest='weightfile')
parser.add_argument('-imgfolder',
type=str,
help='path of your image folder.',
dest='imgfolder')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
if args.imgfolder:
m = Darknet(args.cfgfile)
m.print_network()
m.load_weights(args.weightfile)
print('Loading weights from %s... Done!' % (args.weightfile))
if use_cuda:
m.cuda()
imgfiles = load_images_path(args.imgfolder)
for imgfile in imgfiles:
detect_cv2(m, imgfile)
class Yolov4Node(object):
def __init__(self, cfgfile, weightfile):
rospy.on_shutdown(self.shutdown_cb)
self.model = Darknet(cfgfile)
self.model.print_network()
self.model.load_weights(weightfile)
self.model.eval()
print('Loading weights from %s... Done!' % (weightfile))
self.num_classes = 80
if self.num_classes == 20:
namesfile = os.path.dirname(__file__) + '/data/voc.names'
elif self.num_classes == 80:
namesfile = os.path.dirname(__file__) + '/data/coco.names'
else:
namesfile = os.path.dirname(__file__) + '/data/names'
self.class_names = load_class_names(namesfile)
self.use_cuda = 1
if self.use_cuda:
self.model.cuda()
self.cvbridge = CvBridge()
self.pub_bbox = rospy.Publisher('~det2d_result',
Detection2D,
queue_size=1)
self.sub_image = rospy.Subscriber("/camera/color/image_raw",
ROSImage,
self.image_cb,
queue_size=1)
self.detection_srv = rospy.Service("~yolo_detect", Detection2DTrigger,
self.srv_cb)
print(rospy.get_name() + ' is ready.')
def srv_cb(self, req):
try:
cv_image = self.cvbridge.imgmsg_to_cv2(req.image, "rgb8")
# print("Get image...")
except CvBridgeError as e:
print(e)
return
img_sized = cv2.resize(cv_image, (self.model.width, self.model.height))
boxes_batch = do_detect(self.model, img_sized, 0.5, 0.2, self.use_cuda)
detection_msg = Detection2D()
detection_msg.header.stamp = rospy.Time.now()
detection_msg.header.frame_id = req.image.header.frame_id
# Batch size != 1
if len(boxes_batch) != 1:
print("Batch size != 1, cannot handle it")
exit(-1)
boxes = boxes_batch[0]
# print('num_detections:', len(boxes))
for index, box in enumerate(boxes):
# print('box:', box)
bbox_msg = BBox2D()
bbox_msg.center.x = math.floor(box[0] * req.image.width)
bbox_msg.center.y = math.floor(box[1] * req.image.height)
bbox_msg.size_x = math.floor(box[2] * req.image.width)
bbox_msg.size_y = math.floor(box[3] * req.image.height)
bbox_msg.id = box[6]
bbox_msg.score = box[5]
bbox_msg.class_name = self.class_names[bbox_msg.id]
detection_msg.boxes.append(bbox_msg)
# cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2BGR)
result_img = plot_boxes_cv2(cv_image,
boxes,
savename=None,
class_names=self.class_names,
interest_classes=INTEREST_CLASSES)
detection_msg.result_image = self.cvbridge.cv2_to_imgmsg(
result_img, "bgr8")
print('return {} detection results'.format(len(boxes)))
return Detection2DTriggerResponse(result=detection_msg)
def image_cb(self, msg):
try:
cv_image = self.cvbridge.imgmsg_to_cv2(msg, "rgb8")
rospy.loginfo("Get image")
except CvBridgeError as e:
print(e)
return
img_sized = cv2.resize(cv_image, (self.model.width, self.model.height))
boxes_batch = do_detect(self.model, img_sized, 0.4, 0.3, self.use_cuda)
detection_msg = Detection2D()
detection_msg.header.stamp = rospy.Time.now()
detection_msg.header.frame_id = msg.header.frame_id
detection_msg.result_image = msg
# Batch size != 1
if len(boxes_batch) != 1:
print("Batch size != 1, cannot handle it")
exit(-1)
boxes = boxes_batch[0]
# print('num_detections:', len(boxes))
for index, box in enumerate(boxes):
# print('box:', box)
bbox_msg = BBox2D()
bbox_msg.center.x = math.floor(box[0] * msg.width)
bbox_msg.center.y = math.floor(box[1] * msg.height)
bbox_msg.size_x = math.floor(box[2] * msg.width)
bbox_msg.size_y = math.floor(box[3] * msg.height)
bbox_msg.id = box[6]
bbox_msg.score = box[5]
bbox_msg.class_name = self.class_names[bbox_msg.id]
detection_msg.boxes.append(bbox_msg)
result_img = plot_boxes_cv2(cv_image,
boxes,
savename=None,
class_names=self.class_names,
interest_classes=INTEREST_CLASSES)
detection_msg.result_image = self.cvbridge.cv2_to_imgmsg(
result_img, "bgr8")
self.pub_bbox.publish(detection_msg)
result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)
cv2.imshow('Yolo demo', result_img)
cv2.waitKey(1)
def shutdown_cb(self):
rospy.loginfo("Shutdown " + rospy.get_name())
if hasattr(self, 'model'): del self.model
if hasattr(self, 'cv_bridge'): del self.cv_bridge
if __name__ == '__main__':
dataset = 'patch'
if dataset == 'patch':
from cfg_patch import Cfg
checkpoint = 'checkpoints/Yolov4_modanet_128_epoch46.pth'
elif dataset == 'modanet_whole':
pass
cfg = Cfg
cfg.gpu = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
json_gt = os.path.expanduser('~/data/datasets/modanet/Annots/modanet_instances_val_new.json')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
model.load_model(checkpoint, device)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
#model.load_state_dict(torch.load(checkpoint), strict=True)
val_dataset = YoloModanetHumanDataset(cfg.anno_path, cfg, train=False)
n_val = len(val_dataset)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=True, num_workers=0,
pin_memory=True, drop_last=False, collate_fn=val_collate)
import torch
from cfg import Cfg
from tool.darknet2pytorch import Darknet
from tool.torch_utils import do_detect
from PIL import Image
import cv2
Cfg.cfgfile = '/home/isv/Documents/tensorrt/yolov4/darknet-master/yolov4-tiny.cfg'
Cfg.weights_file = '/home/isv/Documents/tensorrt/yolov4/darknet-master/backup/yolov4-tiny_last.weights'
model = Darknet(Cfg.cfgfile)
model.print_network()
model.load_weights(Cfg.weights_file)
if torch.cuda.is_available():
use_cuda = 1
else:
use_cuda = 0
if use_cuda:
model = model.cuda()
img = cv2.imread('/home/isv/qt_projects/build-final_thread-Desktop_Qt_5_13_2_GCC_64bit-Debug/plateSaved/187.jpg')
img = cv2.resize(img, (320, 320))
# for i in range(10):
boxes = do_detect(model, img, 0.1, 0.4, use_cuda)
numberOfDetection = len(boxes[0])
for i in range(numberOfDetection):
each_box = boxes[0][i]
x = int(320 * each_box[0])
y = int(320 * each_box[1])
w = int(320 * each_box[2])
h = int(320 * each_box[3])
img = cv2.rectangle(img, (x, y), (w, h), (100, 220, 0), 2)
def _get_date_str():
now = datetime.datetime.now()
return now.strftime('%Y-%m-%d_%H-%M')
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
# try:
train(model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device, )
# except KeyboardInterrupt:
# torch.save(model.state_dict(), 'INTERRUPTED.pth')
# logging.info('Saved interrupt')
x, y = int(x1), int(y1)
pred_out = {
"image_id": int(img_id),
"category_id": int(coco_category[category_id]),
# "bbox": [x,y,width,height, int(x1), int(x2), int(y1), int(y2)],
"bbox": [x,y,width,height],
"score": float(conf),
}
coco_pred.append(pred_out)
return coco_pred
if __name__=='__main__':
# Load Model
m = Darknet(opt.cfgfile)
m.load_weights(opt.weights)
use_cuda = opt.use_cuda
if use_cuda:
torch.cuda.set_device(torch.device('cuda:{}'.format(opt.device)))
m.cuda()
# Data Loader
anno = COCO(opt.anno_json)
val_set = COCOImage(opt.anno_json, opt.img_path, opt.img_size)
val_loader = DataLoader(val_set, opt.batch_size, shuffle=True, num_workers=0)
# Accumulate results
result_dict = dict([])
print('STRAT detection')
for imgs, img_ids, sizes in tqdm(val_loader):
class PatchTrainer(object):
def __init__(self, mode):
cfgfile = './cfg/yolov4.cfg'
weightfile = '../common_data/yolov4.weights'
self.darknet_model = Darknet(cfgfile).cuda()
self.darknet_model.print_network()
self.darknet_model.load_weights(weightfile)
self.darknet_model.eval()
print('Loading weights from %s... Done!' % (weightfile))
def train(self):
"""
Optimize a patch to generate an adversarial example.
:return: Nothing
"""
img_size = 800
batch_size = 1
n_epochs = 5000
max_lab = 14
# Dataset prepare
data = CocoTrainPerson(dataType="train2017", num_use=100)
dataloader = DataLoader(data, batch_size=1,
shuffle=False) #使用DataLoader加载数据
ATTACK_AREA_RATE = 0.1
decay_epoch = 100
k = 0
for i_batch, batch_data in enumerate(dataloader):
img, mask, bbox, class_label = batch_data[0][0], batch_data[1][
0], batch_data[2][0], batch_data[3][0]
##############################################################################
img_name = batch_data[4][0]
mask_area = torch.sum(mask)
print('---------------')
print(img_name)
print('---------------')
record_dir = '../common_data/yolo4_dap_attack/disappear/area'
record_path = os.path.join(record_dir,
img_name.split('.')[0] + '.txt')
# use segment SLIC
base_SLIC_seed_num = 3000
img_np = img.numpy().transpose(1, 2, 0)
mask_np = mask.numpy()
numSegments = int(base_SLIC_seed_num / (500 * 500) *
torch.sum(mask))
segments_np = slic(image=img_np,
n_segments=numSegments,
sigma=0,
slic_zero=True,
mask=mask_np)
segments_tensor = torch.from_numpy(segments_np).float().cuda()
segments_label = torch.unique(segments_tensor)
segments_label = segments_label[1:]
zero_layer = torch.zeros_like(segments_tensor)
one_layer = torch.ones_like(segments_tensor)
bbox_x1 = bbox[0]
bbox_y1 = bbox[1]
bbox_w = bbox[2]
bbox_h = bbox[3]
bbox_x_c = bbox_x1 + bbox_w / 2
bbox_y_c = bbox_y1 + bbox_h / 2
bbox_x_c_int = int(bbox_x_c)
bbox_y_c_int = int(bbox_y_c)
# 3 load attack region
load_patch_dir = '../common_data/NES_search_test_1107/' + img_name.split(
'_')[0]
load_patch_list = os.listdir(load_patch_dir)
load_patch_list.sort()
wat_num_max = 0
for i_name in load_patch_list:
wat_num = int(i_name.split('_')[0])
if wat_num > wat_num_max:
wat_num_max = wat_num
for i_name in load_patch_list:
wat_num = int(i_name.split('_')[0])
if wat_num == wat_num_max:
max_name = i_name
break
load_patch = os.path.join(load_patch_dir, max_name)
load_img = Image.open(load_patch).convert('RGB')
load_img = transforms.ToTensor()(load_img)
region_mask = 2 * load_img - img.cpu()
region_mask = torch.sum(region_mask, dim=0) / 3
region_mask = torch.where(mask > 0, region_mask,
torch.zeros_like(region_mask))
attack_region_tmp_pil = transforms.ToPILImage()(region_mask.cpu())
attack_region_tmp_pil.save('013k.png')
# process mask
region_mask_new = torch.zeros_like(region_mask).cuda()
for i in range(segments_label.shape[0]):
sp = segments_label[i]
right_color = (torch.where(segments_tensor == sp,
region_mask.cuda(),
one_layer * (-10))).cpu()
right_color = torch.mean(right_color[right_color != -10])
color_layer = torch.ones_like(segments_tensor).fill_(
right_color)
region_mask_new = torch.where(segments_tensor == sp,
color_layer, region_mask_new)
region_mask_new = region_mask_new
region_mask = region_mask_new
region_mask_unique = torch.unique(region_mask)
for i in range(region_mask_unique.shape[0]):
thres = region_mask_unique[i]
# region_mask_tmp = torch.zeros_like(region_mask)
region_mask_tmp = torch.where(region_mask > thres, one_layer,
zero_layer)
pixel_num = torch.sum(region_mask_tmp)
if pixel_num < mask_area * ATTACK_AREA_RATE:
break
attack_region_search_top = region_mask_tmp
attack_region_search_top = get_conv_envl(attack_region_search_top)
attack_region_tmp_pil = transforms.ToPILImage()(
attack_region_search_top.cpu())
attack_region_tmp_pil.save('012k.png')
attack_region_tmp = attack_region_search_top
attack_region_tmp = attack_region_tmp.cuda()
now_area = float(torch.sum(attack_region_tmp) / mask_area)
print('---------------')
print('You have used ', now_area, 'area.')
print('---------------')
## start at gray
adv_patch_w = torch.zeros(3, 500, 500).cuda()
adv_patch_w.requires_grad_(True)
optimizer = torch.optim.Adam([{
'params': adv_patch_w,
'lr': 0.1
}],
amsgrad=True)
t_op_num = 1500
min_max_iou_record = 1
for t_op_step in range(t_op_num):
adv_patch = torch.sigmoid(adv_patch_w)
patched_img = torch.where(attack_region_tmp > 0, adv_patch,
img.cuda()).unsqueeze(0)
patched_img_rsz = F.interpolate(patched_img, (608, 608),
mode='bilinear').cuda()
yolov4_output = self.darknet_model(patched_img_rsz)
bbox_pred, cls_pred, obj_pred = yolov4_output
bbox_pred = bbox_pred.squeeze()
total_loss = torch.max(obj_pred)
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
person_conf = (cls_pred * obj_pred)[0, :, 0]
ground_truth_bbox = [
bbox_x1, bbox_y1, bbox_x1 + bbox_w, bbox_y1 + bbox_h
]
ground_truth_bbox = torch.Tensor(ground_truth_bbox).unsqueeze(
0).cuda() / 500
ground_truth_bboxs = ground_truth_bbox.repeat(
bbox_pred.shape[0], 1)
iou = compute_iou_tensor(bbox_pred, ground_truth_bboxs)
# ----------------------------------
# ------------------------
# early stop
#test
patched_img_cpu = patched_img.cpu().squeeze()
test_confidence_threshold = 0.45
ov_test_thrs_index = torch.where(
attack_prob > test_confidence_threshold)[0]
final_pbbox = det_bboxes[:,
class_label * 4:(class_label + 1) * 4]
ground_truth_bboxs_final = ground_truth_bbox.repeat(
final_pbbox.shape[0], 1)
iou = compute_iou_tensor(final_pbbox, ground_truth_bboxs_final)
attack_prob_select_by_iou_ = attack_prob[iou > 0.05]
attack_prob_select_by_iou_ = attack_prob_select_by_iou_[
attack_prob_select_by_iou_ > test_confidence_threshold]
# stop if no such class found
if attack_prob_select_by_iou_.shape[0] == 0:
print('Break at', t_op_step, 'no bbox found')
# save image
patched_img_cpu_pil = transforms.ToPILImage()(
patched_img_cpu)
out_file_path = os.path.join(
'../common_data/yolo4_dap_attack/success' +
str(int(ATTACK_AREA_RATE * 100)), img_name)
patched_img_cpu_pil.save(out_file_path)
break
# max same-class object bounding box iou s
final_pbbox = det_bboxes[ov_test_thrs_index][:, class_label *
4:(class_label +
1) * 4]
ground_truth_bboxs_final = ground_truth_bbox.repeat(
final_pbbox.shape[0], 1)
iou = compute_iou_tensor(final_pbbox, ground_truth_bboxs_final)
iou_max = torch.max(iou)
if iou_max < 0.05:
print('Break at', t_op_step, 'iou final max:',
torch.max(iou))
# save image
patched_img_cpu_pil = transforms.ToPILImage()(
patched_img_cpu)
out_file_path = os.path.join(
'../common_data/yolo4_dap_attack/success' +
str(int(ATTACK_AREA_RATE * 100)), img_name)
patched_img_cpu_pil.save(out_file_path)
break
# report
ground_truth_bboxs = ground_truth_bbox.repeat(1000, 1)
final_pbbox = det_bboxes[ov_test_thrs_index][:, class_label *
4:(class_label +
1) * 4]
ground_truth_bboxs_final = ground_truth_bbox.repeat(
final_pbbox.shape[0], 1)
iou = compute_iou_tensor(final_pbbox, ground_truth_bboxs_final)
max_iou = torch.max(iou)
if max_iou < min_max_iou_record:
min_max_iou_record = max_iou
txt_save_dir = '../common_data/yolo4_dap_attack/iou' + str(
int(ATTACK_AREA_RATE * 100))
txt_save_path = os.path.join(
txt_save_dir,
img_name.split('.')[0] + '.txt')
with open(txt_save_path, 'w') as f:
text = str(float(max_iou))
f.write(text)
if t_op_step % 100 == 0:
iou_sort = torch.sort(
iou, descending=True)[0][:6].detach().clone().cpu()
print(t_op_step, 'iou t-cls :', iou_sort)
# iou over 0.5, confidence print
final_pbbox = det_bboxes[:, class_label *
4:(class_label + 1) * 4]
iou = compute_iou_tensor(
final_pbbox,
ground_truth_bbox.repeat(final_pbbox.shape[0], 1))
attack_prob = det_labels[:, class_label]
attack_prob_select_by_iou_ = attack_prob[iou > 0.05]
attack_prob_select_by_iou_sort = torch.sort(
attack_prob_select_by_iou_,
descending=True)[0][:6].detach().cpu()
print(t_op_step, 'right cls cf:',
attack_prob_select_by_iou_sort)
print()
def train(model,
device,
config,
epochs=5,
batch_size=1,
save_cp=True,
log_step=20,
img_scale=0.5):
train_dataset = Yolo_dataset(config.train_label, config, train=True)
val_dataset = Yolo_dataset(config.val_label, config, train=False)
n_train = len(train_dataset)
n_val = len(val_dataset)
train_loader = DataLoader(train_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=collate)
val_loader = DataLoader(val_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=val_collate)
writer = SummaryWriter(
log_dir=config.TRAIN_TENSORBOARD_DIR,
filename_suffix=
f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}',
comment=
f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}'
)
# writer.add_images('legend',
# torch.from_numpy(train_dataset.label2colorlegend2(cfg.DATA_CLASSES).transpose([2, 0, 1])).to(
# device).unsqueeze(0))
max_itr = config.TRAIN_EPOCHS * n_train
# global_step = cfg.TRAIN_MINEPOCH * n_train
global_step = 0
logging.info(f'''Starting training:
Epochs: {epochs}
Batch size: {config.batch}
Subdivisions: {config.subdivisions}
Learning rate: {config.learning_rate}
Training size: {n_train}
Validation size: {n_val}
Checkpoints: {save_cp}
Device: {device.type}
Images size: {config.width}
Optimizer: {config.TRAIN_OPTIMIZER}
Dataset classes: {config.classes}
Train label path:{config.train_label}
Pretrained:
''')
# learning rate setup
def burnin_schedule(i):
if i < config.burn_in:
factor = pow(i / config.burn_in, 4)
elif i < config.steps[0]:
factor = 1.0
elif i < config.steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
if config.TRAIN_OPTIMIZER.lower() == 'adam':
optimizer = optim.Adam(
model.parameters(),
lr=config.learning_rate / config.batch,
betas=(0.9, 0.999),
eps=1e-08,
)
elif config.TRAIN_OPTIMIZER.lower() == 'sgd':
optimizer = optim.SGD(
params=model.parameters(),
lr=config.learning_rate / config.batch,
momentum=config.momentum,
weight_decay=config.decay,
)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
criterion = Yolo_loss(device=device,
batch=config.batch // config.subdivisions,
n_classes=config.classes)
# scheduler = ReduceLROnPlateau(optimizer, mode='max', verbose=True, patience=6, min_lr=1e-7)
# scheduler = CosineAnnealingWarmRestarts(optimizer, 0.001, 1e-6, 20)
save_prefix = 'Yolov4_epoch'
saved_models = deque()
model.train()
for epoch in range(epochs):
# model.train()
epoch_loss = 0
epoch_step = 0
with tqdm(total=n_train,
desc=f'Epoch {epoch + 1}/{epochs}',
unit='img',
ncols=50) as pbar:
for i, batch in enumerate(train_loader):
global_step += 1
epoch_step += 1
images = batch[0]
bboxes = batch[1]
images = images.to(device=device, dtype=torch.float32)
bboxes = bboxes.to(device=device)
bboxes_pred = model(images)
loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = criterion(
bboxes_pred, bboxes)
# loss = loss / config.subdivisions
loss.backward()
epoch_loss += loss.item()
if global_step % config.subdivisions == 0:
optimizer.step()
scheduler.step()
model.zero_grad()
if global_step % (log_step * config.subdivisions) == 0:
writer.add_scalar('train/Loss', loss.item(), global_step)
writer.add_scalar('train/loss_xy', loss_xy.item(),
global_step)
writer.add_scalar('train/loss_wh', loss_wh.item(),
global_step)
writer.add_scalar('train/loss_obj', loss_obj.item(),
global_step)
writer.add_scalar('train/loss_cls', loss_cls.item(),
global_step)
writer.add_scalar('train/loss_l2', loss_l2.item(),
global_step)
writer.add_scalar('lr',
scheduler.get_lr()[0] * config.batch,
global_step)
pbar.set_postfix(
**{
'loss (batch)': loss.item(),
'loss_xy': loss_xy.item(),
'loss_wh': loss_wh.item(),
'loss_obj': loss_obj.item(),
'loss_cls': loss_cls.item(),
'loss_l2': loss_l2.item(),
'lr': scheduler.get_lr()[0] * config.batch
})
logging.debug(
'Train step_{}: loss : {},loss xy : {},loss wh : {},'
'loss obj : {}，loss cls : {},loss l2 : {},lr : {}'.
format(global_step, loss.item(), loss_xy.item(),
loss_wh.item(), loss_obj.item(),
loss_cls.item(), loss_l2.item(),
scheduler.get_lr()[0] * config.batch))
pbar.update(images.shape[0])
if cfg.use_darknet_cfg:
eval_model = Darknet(cfg.cfgfile, inference=True)
else:
eval_model = Yolov4(cfg.pretrained,
n_classes=cfg.classes,
inference=True)
# eval_model = Yolov4(yolov4conv137weight=None, n_classes=config.classes, inference=True)
if torch.cuda.device_count() > 1:
eval_model.load_state_dict(model.module.state_dict())
else:
eval_model.load_state_dict(model.state_dict())
eval_model.to(device)
evaluator = evaluate(eval_model, val_loader, config, device)
del eval_model
stats = evaluator.coco_eval['bbox'].stats
writer.add_scalar('train/AP', stats[0], global_step)
writer.add_scalar('train/AP50', stats[1], global_step)
writer.add_scalar('train/AP75', stats[2], global_step)
writer.add_scalar('train/AP_small', stats[3], global_step)
writer.add_scalar('train/AP_medium', stats[4], global_step)
writer.add_scalar('train/AP_large', stats[5], global_step)
writer.add_scalar('train/AR1', stats[6], global_step)
writer.add_scalar('train/AR10', stats[7], global_step)
writer.add_scalar('train/AR100', stats[8], global_step)
writer.add_scalar('train/AR_small', stats[9], global_step)
writer.add_scalar('train/AR_medium', stats[10], global_step)
writer.add_scalar('train/AR_large', stats[11], global_step)
if save_cp:
try:
# os.mkdir(config.checkpoints)
os.makedirs(config.checkpoints, exist_ok=True)
logging.info('Created checkpoint directory')
except OSError:
pass
save_path = os.path.join(config.checkpoints,
f'{save_prefix}{epoch + 1}.pth')
torch.save(model.state_dict(), save_path)
logging.info(f'Checkpoint {epoch + 1} saved !')
saved_models.append(save_path)
if len(saved_models) > config.keep_checkpoint_max > 0:
model_to_remove = saved_models.popleft()
try:
os.remove(model_to_remove)
except:
logging.info(f'failed to remove {model_to_remove}')
writer.close()
formatter = logging.Formatter(fmt)
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)
return logging
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.weight, n_classes=cfg.classes)
# pretrained_dict = torch.load(cfg.load)
# model_dict = model.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if
# k in model_dict} # pretrained_dict只保留了model_dict中存在的键。为什么直接load它会报错。要先给model_dict更新。
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
total_params = sum(p.numel() for p in model.parameters())
print(f'{total_params:,} total parameters.')
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print(f'{total_trainable_params:,} training parameters.')
def _get_date_str():
now = datetime.datetime.now()
return now.strftime('%Y-%m-%d_%H-%M')
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
if cfg.pretrained:
model.load_weights(cfg.pretrained)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
if cfg.load:
pretrained_dict = torch.load(cfg.load,
map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
console.setLevel(log_level)
formatter = logging.Formatter(fmt)
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)
return logging
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
model = Darknet(cfg.model_config)
model.print_network()
model.load_weights(cfg.weights_file)
model.eval() # set model away from training
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
annotations_file_path = cfg.gt_annotations_path
with open(annotations_file_path) as annotations_file:
try:
annotations = json.load(annotations_file)
except:
def _get_date_str():
now = datetime.datetime.now()
return now.strftime('%Y-%m-%d_%H-%M')
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(
model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device,
)
except KeyboardInterrupt:
from tool.torch_utils import do_detect
from PIL import Image
import matplotlib.pyplot as plt
if __name__ == '__main__':
cfgfile = "./cfg/yolov4.cfg"
weightfile = "./pretrained/yolov4.weights"
namesfile = 'data/coco.names'
imgfile = "./image/image/init.png"
imgfile2 = "./media/059_L.png"
use_cuda = True
# Initiate model
t0 = time.time()
darknet_model = Darknet(cfgfile)
width, height = (darknet_model.width, darknet_model.height)
darknet_model.load_weights(weightfile)
if use_cuda:
darknet_model.cuda()
class_names = load_class_names(namesfile)
t1 = time.time()
total_time = round(t1 - t0, 2)
print("1 - Initiated DepthModel. -- {} minutes {} seconds".format(
total_time // 60, total_time % 60))
print("====================================")
print("====================================")
print("====================================")
# Inference
return parser.parse_args()
if __name__ == '__main__':
cfgfile = "cfg/yolov4.cfg"
weightsfile = "weight/yolov4.weights"
args = arg_parse()
confidence = float(args.confidence)
nms_thesh = float(args.nms_thresh)
CUDA = torch.cuda.is_available()
num_classes = 80
bbox_attrs = 5 + num_classes
class_names = load_class_names("data/coco.names")
model = Darknet(cfgfile)
model.load_weights(weightsfile)
if CUDA:
model.cuda()
model.eval()
cap = cv2.VideoCapture(0)
assert cap.isOpened(), 'Cannot capture source'
frames = 0
start = time.time()
while cap.isOpened():
ret, frame = cap.read()
if ret:
"""hyper parameters"""
use_cuda = True
cfgfile_RGB='./cfg/yolo-1031_obj.cfg'
weightfile_RGB='./weight/yolo-1031_obj_final.weights'
cfgfile_TRM='./cfg/yolo-1031_obj.cfg'
weightfile_TRM='./weight/yolo-1031_trm_obj_final.weights'
imgfile_RGB='./data/RGB/RGB_costco4.avi'
imgfile_TRM='./data/TRM/Thermal_costco4.avi'
imgfile_fusion='./data/fusion/fusion_costco4.avi'
import cv2
RGB = Darknet(cfgfile_RGB)
RGB.print_network()
RGB.load_weights(weightfile_RGB)
print('Loading RGB weights from %s... Done!' % (weightfile_RGB))
if use_cuda:
RGB.cuda()
num_classes = RGB.num_classes
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/1104_obj.names'
class_names = load_class_names(namesfile)
