def __init__(self):
print(os.getcwd())
self.model_cfg = "./src/akhenaten_dv/scripts/Perception/BBoxDetection/model_cfg/yolo_baseline_tiny.cfg"
self.weights_path = './src/akhenaten_dv/scripts/Perception/BBoxDetection/7.weights'
self.conf_thres = 0.8
self.nms_thres = 0.25
self.vanilla_anchor = False
self.xy_loss = 2
self.wh_loss = 1.6
self.no_object_loss = 25
self.object_loss = 0.1
cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if cuda else 'cpu')
random.seed(0)
torch.manual_seed(0)
if cuda:
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = True
torch.cuda.empty_cache()
self.model = Darknet(config_path=self.model_cfg,
xy_loss=self.xy_loss,
wh_loss=self.wh_loss,
no_object_loss=self.no_object_loss,
object_loss=self.object_loss,
vanilla_anchor=self.vanilla_anchor)
# Load weights
self.model.load_weights(self.weights_path,
self.model.get_start_weight_dim())
self.model.to(self.device, non_blocking=True)
def __init__(self, bgr=True, gpu_device=0, **kwargs):
self.__dict__.update(self._defaults) # set up default values
# for portability between keras-yolo3/yolo.py and this
if 'model_path' in kwargs:
kwargs['weights'] = kwargs['model_path']
if 'score' in kwargs:
kwargs['thresh'] = kwargs['score']
self.__dict__.update(kwargs) # update with user overrides
self.class_names = self._get_class()
self.model = Darknet(self.config)
self.model.load_darknet_weights(self.weights)
self.device = gpu_device
self.model.cuda(self.device)
self.model.eval()
self.bgr = bgr
if self.half:
self.model.half()
# warm up
self._detect([np.zeros((10, 10, 3), dtype=np.uint8)])
print('Warmed up!')
def backbone(self, config_path) -> Darknet: # 删除config_path参数
if 'darknet53' in dir(self) and self.darknet53 is not None:
raise Exception(
"Call 'backbone' before 'features' and 'multiscale_features' method")
if self._models_dir is not None and self._pretrained:
self.darknet53 = Darknet(config_path)
self.darknet53.load_state_dict(torch.load(
os.path.join(self._models_dir, 'darknet53-5d3b4d8f.pth')
))
else:
self.darknet53 = Darknet(config_path)
# list(darknet53.children()) consists of following modules
# [0] = Conv2d, [1] = BatchNorm2d, [2] = ReLU,
# [3] = MaxPool2d, [4] = Sequential(Bottleneck...),
# [5] = Sequential(Bottleneck...),
# [6] = Sequential(Bottleneck...),
# [7] = Sequential(Bottleneck...),
# [8] = AvgPool2d, [9] = Linear
# children = list(self.darknet53.children())
# self.features = children[:-3]
# self.num_features_out = 1024
# self.hidden = children[-3]
# self.num_hidden_out = 2048
return self.darknet53
def test(img_path: str = 'data/custom/images/', anno_path: str = 'data/custom/annos/') -> None:
# transform something
initialize(img_path=img_path, anno_path=anno_path, split_ratio=1.0)
# some common config
iou_thres = 0.5
conf_thres = 0.01
nms_thres = 0.5
img_size = 416
batch_size = 16
device = tc.device('cuda' if tc.cuda.is_available() else 'cpu')
# other paths
weights_path = 'my/yolov3_ckpt_1.pth'
model_def = 'config/custom.cfg'
test_path = 'data/custom/train.txt'
class_path = 'data/custom/classes.names'
# load model
class_names = load_classes(class_path)
model = Darknet(model_def).to(device)
model.load_state_dict(tc.load(weights_path))
imgs, img_detections = detect(
model=model,
path=img_path,
conf_thres=conf_thres,
nms_thres=nms_thres,
img_size=img_size,
batch_size=batch_size,
n_cpu=8,
device=device,
)
os.makedirs('predicted_file', exist_ok=True)
class1 = open('predicted_file/det_test_core.txt', 'w')
class2 = open('predicted_file/det_test_coreless.txt', 'w')
for path, boxes in zip(imgs, img_detections):
w, h = Image.open(path).size
boxes = rescale_boxes(boxes, img_size, (h, w))
for box in boxes:
line = [
# os.path.split(path)[1].split('_')[1].split('.')[0], # no prefix
os.path.split(path)[1].split('.')[0], # with prefix 'core_' or 'coreless_'
f'{box[4].tolist():.3f}', # conf
f'{box[0].tolist():.1f}',
f'{box[1].tolist():.1f}',
f'{box[2].tolist():.1f}',
f'{box[3].tolist():.1f}',
]
if box[-1] == 0.0:
class1.write(' '.join(line) + '\n')
elif box[-1] == 1.0:
class2.write(' '.join(line) + '\n')
class1.close()
class2.close()
print('Output file saved.\n')
def main():
args = create_prune_argparser()
config = create_config(args)
# Initialize
init_seeds(seed=0)
model = Darknet(cfg=config['cfg'], arc=config['arc'])
mask = create_mask(model)
bckp = create_backup(model)
device = select_device(config['device'])
model = model.to(device)
# print('Making forwards by 100 iterations')
# mask = mask.to(device)
# x = torch.Tensor(10, 3, 416, 416).to(device)
# for i in range(100):
# out = model(x)
# exit()
data_dict = parse_data_cfg(config['data'])
train_path = data_dict['train']
dataset = LoadImagesAndLabels(
path=train_path,
img_size=config['img_size'][0],
batch_size=config['batch_size'],
augment=True,
hyp=config['hyp'],
cache_images=config['cache_images'],
)
# Dataloader
nw = min([os.cpu_count(), 18 if 18 > 1 else 0, 8]) # number of workers
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=18,
num_workers=nw,
pin_memory=True,
collate_fn=dataset.collate_fn)
# torch.cuda.empty_cache()
imgs, _, _, _ = next(iter(dataloader))
imgs = imgs.float() / 255.0
imgs = imgs.to(device)
start = datetime.datetime.now()
print(f'Starting to compute the time at {start}')
for i in range(10):
prune_on_cpu(model, mask, bckp, imgs, config, device)
end = datetime.datetime.now()
print(f'Ending at {end}')
result = end - start
print(f'Time of {result}')
def run_embedding(aligned_face, opt):
device = "cuda" if opt.use_cuda and torch.cuda.is_available() else "cpu"
model = Darknet(opt.config_path, img_size=160)
model.load_weights(opt.weights_path)
model = model.to(device)
img = np.array(aligned_face)
img = np.repeat(img[np.newaxis, :, :], 3, axis=0)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float().to(device)
embedding = model(img)
embedding = torch.norm(embedding).detach().cpu()
return embedding
def __init__(self, device, model_def, load_path, reg_threshold,
cls_threshold, nms_threshold, image_size):
self.image_size = image_size
self.model = Darknet(model_def, img_size=self.image_size).to(device)
# TODO
# change device to GPU
self.model.load_state_dict(torch.load(load_path, map_location='cpu'))
self.model.eval()
self.reg_threshold = reg_threshold
self.cls_threshold = cls_threshold
self.nms_threshold = nms_threshold
self.device = device
def main():
device = torch.device("cuda" if OPT.use_cuda else "cpu")
cudnn.benchmark = True
classes = tools.load_classes(OPT.class_path) # Extracts class labels from file
model = Darknet(OPT.config_path, img_size=OPT.img_size)
model.load_weights(OPT.weights_path)
if OPT.source == 'image':
detect_image(model, device, classes)
elif OPT.source == 'video':
detect_video(model, device, classes)
def run_compare(cfg, data, prune_cfg, batch_size, origin_weights):
device = select_device('', apex=None, batch_size=batch_size)
if device.type != 'cpu' and torch.cuda.device_count() > 1:
dist.init_process_group(
backend='nccl', # 'distributed backend'
init_method=
'tcp://127.0.0.1:9999', # distributed training init method
world_size=1, # number of nodes for distributed training
rank=0) # distributed training node rank
init_seeds()
data_dict = parse_data_cfg(data)
train_path = data_dict['valid']
dataset = LoadImagesAndLabels(
train_path,
416,
batch_size,
augment=True,
hyp=hyp, # augmentation hyperparameters
rect=False, # rectangular training
cache_labels=True,
cache_images=False)
batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
train_loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=nw,
shuffle=True, # Shuffle=True unless rectangular training is used
pin_memory=True,
collate_fn=dataset.collate_fn)
origin_model = Darknet(cfg).to(device)
chkpt = torch.load(origin_weights, map_location=device)
origin_model.load_state_dict(chkpt['model'], strict=True)
aux_util = AuxNetUtils(origin_model, hyp)
del chkpt
for layer in aux_util.pruning_layer[1:]:
# greedy_channel_select(origin_model, prune_cfg, origin_weights, layer, device, aux_util, train_loader, 0.75)
random_greedy_channel_select(origin_model, prune_cfg, origin_weights,
layer, device, aux_util, train_loader,
0.75)
def transform_to_onnx(cfg_file, weight_file, batch_size, in_h, in_w):
model = Darknet(cfg_file)
pre_dict = torch.load(weight_file, map_location=torch.device('cpu'))
model.load_state_dict(pre_dict['model'])
x = torch.ones((batch_size, 3, in_h, in_w), requires_grad=True)*120 /255.0
onnx_file_name = 'model/yolov3.onnx'
torch.onnx.export(model,
x,
onnx_file_name,
#export_params=True,
#opset_version=11,
#do_constant_folding=True,
input_names=['input'], output_names=['output1'])
#dynamic_axes=None)
print('Onnx model exporting done')
return onnx_file_name, x
def main(target_path, output_path, weights_path, model_cfg, conf_thres,
nms_thres, xy_loss, wh_loss, no_object_loss, object_loss,
vanilla_anchor):
cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
random.seed(0)
torch.manual_seed(0)
if cuda:
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = True
torch.cuda.empty_cache()
model = Darknet(config_path=model_cfg,
xy_loss=xy_loss,
wh_loss=wh_loss,
no_object_loss=no_object_loss,
object_loss=object_loss,
vanilla_anchor=vanilla_anchor)
# Load weights
model.load_weights(weights_path, model.get_start_weight_dim())
model.to(device, non_blocking=True)
detect(target_path,
output_path,
model,
device=device,
conf_thres=conf_thres,
nms_thres=nms_thres)
def yolov3(linux):
# Code
if linux:
YOLO_PATH = '/home/student/Desktop/Automated-Active-Surveillance-System-in-the-Detection-of-Cold-Steel-Weapons/dl_models/yolov3'
else:
base_path = os.path.dirname(os.path.realpath(__file__)) + '\\'
YOLO_PATH = base_path + 'yolov3'
sys.path.insert(0, YOLO_PATH)
from models import Darknet
opt = YoloConfig(linux)
img_size = opt.img_size
out, source, weights, half, view_img, save_txt = opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Darknet(opt.cfg, img_size)
model.load_state_dict(torch.load(weights)['model'])
return model, opt
def main():
img_size = 512 # 必须是32的整数倍 [416, 512, 608]
cfg = "cfg/yolov3-spp.cfg"
weights = "weights/yolov3-spp-ultralytics-{}.pt".format(img_size)
img_path = "test.jpg"
input_size = (img_size, img_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Darknet(cfg, img_size)
model.load_state_dict(torch.load(weights, map_location=device)["model"])
model.to(device)
model.eval()
# init
img = torch.zeros((1, 3, img_size, img_size), device=device)
model(img)
img_o = cv2.imread(img_path) # BGR
assert img_o is not None, "Image Not Found " + img_path
img = img_utils.letterbox(img_o, new_shape=input_size, auto=True, color=(0, 0, 0))[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device).float()
img /= 255.0 # scale (0, 255) to (0, 1)
img = img.unsqueeze(0) # add batch dimension
t1 = torch_utils.time_synchronized()
pred = model(img)[0] # only get inference result
t2 = torch_utils.time_synchronized()
print(t2 - t1)
pred = utils.non_max_suppression(pred, conf_thres=0.3, iou_thres=0.6, multi_label=True)[0]
t3 = time.time()
print(t3 - t2)
# process detections
pred[:, :4] = utils.scale_coords(img.shape[2:], pred[:, :4], img_o.shape).round()
print(pred.shape)
bboxes = pred[:, :4].detach().cpu().numpy()
scores = pred[:, 4].detach().cpu().numpy()
classes = pred[:, 5].detach().cpu().numpy().astype(np.int) + 1
category_index = dict([(i + 1, str(i + 1)) for i in range(90)])
img_o = draw_box(img_o[:, :, ::-1], bboxes, classes, scores, category_index)
plt.imshow(img_o)
plt.show()
img_o.save("test_result.jpg")
def eval_flowchart(init_style, para_part, reg, alpha, ablation_type):
'''Main body for evaluation'''
args = all_args()
# Storage path 'eval/'
os.makedirs(args.save_folder, exist_ok=True)
# Dataset
dataset = get_data(args)
# Load net
net = Darknet(args.model_config_path, img_size=args.img_size)
# Visdom
viz, epoch_aps = init_viz(args, init_style, para_part, reg, alpha, dataset)
# Evaluate
ckpt_path, ckpts = get_ckpts(args, init_style, para_part, reg, alpha,
ablation_type)
mAP_max = 0
for ckpt_idx, ckpt in enumerate(ckpts):
# sample one for hyperparameter adjustment
if ckpt_idx < 120:
continue
# Make output dir
dir_name = '_'.join([
ablation_type, args.arc, args.dataset, args.set_type, init_style,
para_part, reg,
str(alpha), ckpt,
str(ckpt_idx)
])
output_dir = get_output_dir(args.save_folder, dir_name)
# Load weight
args.weight_path = os.path.join(ckpt_path, ckpt)
# assert os.path.isfile(args.weight_path)
try:
net.load_weights(args.weight_path)
except FileNotFoundError as err:
print(err)
# Cuda or not
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
net.eval()
print('Finished loading model!')
# Evaluation, use_07_eval False
aps, mAP = test_net(output_dir,
net,
args.cuda,
dataset,
args.score_thres,
args.nms_thres,
use_07_eval=False,
iou_thres=args.iou_thres)
# If not greater than before, delete
if mAP_max >= mAP:
rmtree(output_dir)
else:
mAP_max = mAP
# Visdom
update_vis(viz, epoch_aps, ckpt_idx + 1, *aps, mAP)
def __init__(self, data_config, net_config, weights, image_size):
"""Class init function."""
QtCore.QThread.__init__(self)
self.image_list = []
self.threshold = 0.9
self.image_directory = ''
self.data = None
self.image_size = image_size
if torch.cuda.is_available():
self.device = torch.device('cuda:0')
else:
self.device = torch.device('cpu')
self.data_config = parse_data_config(data_config)
# Extracts class labels from file
self.classes = load_classes(self.data_config['names'])
self.model = Darknet(net_config, image_size)
checkpoint = torch.load(weights, map_location='cpu')
self.model.load_state_dict(checkpoint['model'])
def build_model(self) -> nn.Module:
opt = get_cli_args(batch_size=pedl_batch_size,
prebias=pedl_prebias,
accumulate=pedl_accumulate)
hyp = get_hyp(lr0=pedl_init_lr)
# Initialize model
model = Darknet(opt.cfg, arc=opt.arc) # .to(device)
# Fetch starting weights
# TODO Once download_data_fn is implemented this should go into download_data
attempt_download(opt.weights)
chkpt = torch.load(opt.weights)
# load model
try:
chkpt["model"] = {
k: v
for k, v in chkpt["model"].items()
if model.state_dict()[k].numel() == v.numel()
}
model.load_state_dict(chkpt["model"], strict=False)
except KeyError as e:
s = (
"%s is not compatible with %s. Specify --weights '' or specify a --cfg compatible with %s. "
"See https://github.com/ultralytics/yolov3/issues/657" %
(opt.weights, opt.cfg, opt.weights))
raise KeyError(s) from e
del chkpt
data_dict = get_data_cfg()
nc = 1 if opt.single_cls else int(data_dict["classes"])
model.nc = nc # attach number of classes to model
model.arc = opt.arc # attach yolo architecture
model.hyp = hyp # attach hyperparameters to model
train_dataset = LazyModule.get()
# The model class weights depend on the dataset labels
model.class_weights = labels_to_class_weights(
train_dataset.labels, nc) # attach class weights
return model
class DarkNet53(backbone.base.Base):
def __init__(self, pretrained: bool, models_dir: Optional[str] = None):
super().__init__(pretrained)
self._models_dir = models_dir
def features(self) -> Tuple[nn.Module, nn.Module, int, int]:
if 'darknet53' not in dir(self):
self.backbone()
return self.darknet53
def backbone(self, config_path) -> Darknet: # 删除config_path参数
if 'darknet53' in dir(self) and self.darknet53 is not None:
raise Exception(
"Call 'backbone' before 'features' and 'multiscale_features' method")
if self._models_dir is not None and self._pretrained:
self.darknet53 = Darknet(config_path)
self.darknet53.load_state_dict(torch.load(
os.path.join(self._models_dir, 'darknet53-5d3b4d8f.pth')
))
else:
self.darknet53 = Darknet(config_path)
# list(darknet53.children()) consists of following modules
# [0] = Conv2d, [1] = BatchNorm2d, [2] = ReLU,
# [3] = MaxPool2d, [4] = Sequential(Bottleneck...),
# [5] = Sequential(Bottleneck...),
# [6] = Sequential(Bottleneck...),
# [7] = Sequential(Bottleneck...),
# [8] = AvgPool2d, [9] = Linear
# children = list(self.darknet53.children())
# self.features = children[:-3]
# self.num_features_out = 1024
# self.hidden = children[-3]
# self.num_hidden_out = 2048
return self.darknet53
def convert(cfg='cfg/yolov4.cfg', weights='weights/yolov4.weights'):
# Initialize model
model = Darknet(cfg)
if weights.endswith('.weights'): # darknet format
_ = load_darknet_weights(model, weights)
chkpt = {
'epoch': -1,
'best_fitness': None,
'training_results': None,
'model': model.state_dict(),
'optimizer': None
}
torch.save(chkpt, 'weights/yolov4.pt')
print("Success: converted '%s' to 'converted.pt'" % weights)
else:
print('Error: extension not supported.')
class Detector(object):
def __init__(self, device, model_def, load_path, reg_threshold,
cls_threshold, nms_threshold, image_size):
self.image_size = image_size
self.model = Darknet(model_def, img_size=self.image_size).to(device)
# TODO
# change device to GPU
self.model.load_state_dict(torch.load(load_path, map_location='cpu'))
self.model.eval()
self.reg_threshold = reg_threshold
self.cls_threshold = cls_threshold
self.nms_threshold = nms_threshold
self.device = device
@torch.no_grad()
def __call__(self, image):
original_size = image.shape[:2]
tensor = torch.from_numpy(image).to(self.device).permute(2, 0, 1)
tensor = tensor.contiguous().float().div_(255)
_, h, w = tensor.shape
dim_diff = np.abs(h - w)
pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
tensor = f.pad(tensor, pad, "constant", value=0)
tensor = f.interpolate(tensor.unsqueeze(0),
size=self.image_size,
mode="nearest").squeeze_(0)
result = self.model(tensor.unsqueeze_(0))
detection = non_max_suppression(result, self.reg_threshold,
self.nms_threshold)[0]
if detection is not None:
detection = detection[detection[:, -2] > self.cls_threshold]
detection = rescale_boxes(detection, self.image_size,
original_size)
else:
print("detection result is None")
return detection
def main(*, batch_size, model_cfg, weights_path, bbox_all, step, n_cpu):
cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
random.seed(0)
torch.manual_seed(0)
if cuda:
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = True
torch.cuda.empty_cache()
# Initiate model
model = Darknet(model_cfg)
validate_uri, _, weights_uri = model.get_links()
_, _, _, _, bw = model.get_dataAug()
num_images, _ = model.num_images()
# Load weights
model.load_weights(weights_path, model.get_start_weight_dim())
model.to(device, non_blocking=True)
# Get dataloader
dataloader = torch.utils.data.DataLoader(ImageLabelDataset(
validate_uri,
height=img_height,
width=img_width,
augment_hsv=False,
augment_affine=False,
num_images=num_images,
bw=bw,
n_cpu=n_cpu,
lr_flip=False,
ud_flip=False),
batch_size=batch_size,
shuffle=False,
num_workers=n_cpu,
pin_memory=True)
return validate(dataloader, model, device, step, bbox_all)
def main(args):
import torch
from os.path import join, dirname
from models import Darknet, save_weights, parse_model_cfg
from utils.prune_utils import write_cfg
cfg = parse_model_cfg(args.cfg)
idx = list(range(5))
if args.class_idx != '-1':
idx += [int(x) + 5
for x in args.class_idx.split(',')] # box + obj + cls
else:
for i, c in enumerate(cfg):
if c['type'] == 'yolo':
nf = int(cfg[i - 1]['filters'])
nc = int(c['classes'])
na = nf // (5 + nc)
c['classes'] = '0'
cfg[i - 1]['filters'] = str(nf - na * nc)
model = Darknet(cfg)
na = model.module_list[model.yolo_layers[0]].na # anchor number per point
sd0 = model.state_dict()
if args.src.endswith('.pt'):
sd1 = torch.load(args.src)['model']
else:
raise NotImplementedError('darknet weights not supported')
for k in sd0:
if sd0[k].shape != sd1[k].shape:
sd1[k] = torch.cat([x[idx] for x in sd1[k].chunk(na)])
model.load_state_dict(sd1)
save_weights(model, args.dst)
from os.path import basename, splitext
cfg_path = args.cfg.replace(
splitext(basename(args.cfg))[0],
splitext(basename(args.dst))[0])
write_cfg(cfg_path, [model.hyperparams] + model.module_defs)
print('cfg saved to:', cfg_path)
def Load_Yolo(device):
#Load Darknet
yolo_model_def = os.path.join(yolo_path, 'config/yolov3-tiny.cfg')
yolo_img_size = 416
yolo_weights_path = os.path.join(yolo_path, 'weights/yolov3-tiny.weights')
model = Darknet(yolo_model_def, img_size=yolo_img_size).to(device)
if yolo_weights_path.endswith(".weights"):
# Load darknet weights
model.load_darknet_weights(yolo_weights_path)
else:
# Load checkpoint weights
model.load_state_dict(torch.load(yolo_weights_path))
model.eval() # Set in evaluation mode
return model
def load_checkpoint(file_name, device, S, B, C, cfg=None):
print('Loading checkpoint...')
if cfg is None:
checkpoint = torch.load(file_name)
cfg = checkpoint['cfg']
BACKBONE = cfg['MODEL']
N_LAYERS = cfg.get('N_LAYERS', 0)
MIN_IMAGES = cfg['MIN_IMAGES']
DATASET_DIR = cfg['DATASET_DIR']
print(MIN_IMAGES, DATASET_DIR)
print('backbone:', BACKBONE)
if BACKBONE == 'Darknet':
n_classes = get_n_classes(MIN_IMAGES, root=DATASET_DIR)
backbone = Darknet(n_layers=N_LAYERS, num_classes=n_classes)
backbone.load_state_dict(checkpoint['model'])
model = YoloV1_pretrained(backbone=backbone,
n_layers=N_LAYERS,
grid_size=S,
num_boxes=B,
num_classes=C).to(device)
# for param in backbone.parameters():
# param.requires_grad = False
elif BACKBONE == 'VGG16':
model = OD_backbone(bb_name=BACKBONE,
grid_size=S,
num_boxes=B,
num_classes=C).to(device)
if BACKBONE == 'Darknet':
epoch = checkpoint['epoch']
else:
epoch = 0
return model, cfg, epoch
def yolov3(linux):
# image_path = "/home/student/Desktop/Automated-Active-Surveillance-System-in-the-Detection-of-Cold-Steel-Weapons/models/yolov3/meme.jpg"
# # os.system(f"cd yolov3")
# os.chdir("yolov3")
# os.system(f"python3 detect.py --source {image_path} --cfg cfg/yolov3-spp.cfg --weights yolov3-spp.pt")
# os.chdir("..")
# YOLOV3_PATH = '/home/student/Desktop/Automated-Active-Surveillance-System-in-the-Detection-of-Cold-Steel-Weapons/models/yolov3'
# sys.path.insert(0, YOLOV3_PATH)
# Code
if linux:
YOLO_PATH = '/home/student/Desktop/Automated-Active-Surveillance-System-in-the-Detection-of-Cold-Steel-Weapons/dl_models/yolov3'
else:
YOLO_PATH = 'D:\GitLab_respos\Automated-Active-Surveillance-System-in-the-Detection-of-Cold-Steel-Weapons\dl_models\yolov3'
sys.path.insert(0, YOLO_PATH)
from models import Darknet
opt = YoloConfig(linux)
img_size = opt.img_size
out, source, weights, half, view_img, save_txt = opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Darknet(opt.cfg, img_size)
model.load_state_dict(torch.load(weights)['model'])
return model, opt
def restore_pruned_checkpt(cfg, img_size, resume_prune_point):
device = torch_utils.select_device()
model = Darknet(cfg, img_size)
pruned_weight = resume_prune_point + os.sep + pruned_pt_name
pruned_tracker = resume_prune_point + os.sep + pruned_tracker_name
load_tracker(pruned_tracker, model)
checkpoint = torch.load(pruned_weight, map_location='cpu')
model.load_state_dict(checkpoint['model'])
model.to(device)
return model, pruned_tracker
def initial_model_gen(cfg, img_size, weights):
device = torch_utils.select_device()
model = Darknet(cfg, img_size)
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
else: # darknet format
assert False, 'initial point needs to be pytorch format'
model.to(device).eval()
tracker = PruneTracker()
model(torch.rand(1, 3, img_size, img_size).to(device), tracker=tracker)
return model, tracker
def create_model(opt):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Darknet(opt.cfg).to(device)
#model.apply(weights_init_normal)
if opt.pretrained_weights:
if opt.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(opt.pretrained_weights))
else:
model.load_darknet_weights(opt.pretrained_weights)
model_cfg = parse_model_config(opt.cfg)
return model, model_cfg
def setup_detector(opt):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
weights_path = os.path.join(
opt.weights_path,
"weights_RADAR.pth" if opt.radar else "weights_LIDAR.pth")
# Set up model
model = Darknet(opt.model_def, img_size=cnf.BEV_WIDTH).to(device)
# Load checkpoint weights
model.load_state_dict(torch.load(weights_path, map_location=device))
# Eval mode
model.eval()
return model
def main(image_uri,
output_uri,
weights_uri,
model_cfg,
img_size,
bw,
conf_thres,
nms_thres):
cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
random.seed(0)
torch.manual_seed(0)
if cuda:
torch.cuda.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = True
torch.cuda.empty_cache()
model = Darknet(model_cfg, img_size)
# Load weights
weights_path = storage_client.get_file(weights_uri)
if weights_path.endswith('.weights'): # darknet format
model.load_weights(weights_path)
elif weights_path.endswith('.pt'): # pytorch format
checkpoint = torch.load(weights_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
model.to(device, non_blocking=True)
detect(image_uri,
output_uri,
model,
img_size,
bw,
device=device,
conf_thres=conf_thres,
nms_thres=nms_thres)
def main():
# model
model = Darknet(str(args.config), img_size=416)
model_wts = torch.load(str(args.weight), map_location='cpu')
model.load_state_dict(model_wts)
if torch.cuda.is_available():
print('gpu: available')
model = model.cuda()
else:
print('gpu: not available')
# read video
print(">> reading video...")
video, info = read_video(args.input)
video = video[:, :, :, ::-1]
print(" shape:", video.shape)
print(" info: ", info)
# forward
print(">> predicting...")
imgs, bboxes, ss, labels = [], [], [], []
for i in tqdm(range(0, len(video))):
img = video[i]
bbox, max_s = model.predict(img, args.conf_thresh, args.nms_thresh)
imgs.append(img)
ss.append(max_s)
if len(bbox) != 0:
bboxes.append([bbox])
labels.append([0])
else:
bboxes.append([])
labels.append([])
# draw bbox
imgs = np.asarray(imgs)
# imgs = imgs[:,:,::-1]
for i in tqdm(range(len(imgs))):
if len(bboxes[i]) != 0:
ty, tx, by, bx = [int(n) for n in bboxes[i][0]]
cv2.rectangle(imgs[i], (tx, ty), (bx, by), (255, 0, 0), 3)
# save as video
print(">> saving video...")
imgs = imgs[:, :, :, ::-1]
write_video(args.output, imgs, info)
