def main(args):
cfg = args.cfg
obj_data = args.obj_data
weights = args.weights
img_path = args.img_path
# create save dir and anno file
labels = {"pen", "box", "bottle", "spoon"}
create_img_dir(labels)
create_anno_file()
# read an image.
frame = cv2.imread(img_path)
# setup yolo model
yolo = YOLO()
net = yolo.load_net(cfg.encode(), weights.encode(), 0)
meta = yolo.load_meta(obj_data.encode())
# inference
r = yolo.detect(net, meta, img_path.encode())
# post process
for i, r0 in enumerate(r):
bbox_info = get_bbox(r0)
bbox_img, label = get_bbox_img(frame, bbox_info)
save_img(bbox_img, label, i)
cv2.imshow("bbox img", bbox_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
def main():
set_random_seed(0)
device = 'cpu'
model = YOLO(name='YOLOv2', model='models/yolov2-ss.cfg', device=device)
dataset = PascalDatasetYOLO(root_dir=['../../../Data/SS/'],
class_file='../../../Data/SS/ss.names',
dataset=['test'],
batch_size=32,
image_size=model.default_image_size,
anchors=model.anchors,
strides=model.strides,
skip_difficult=False,
do_transforms=False,
multi_scale=model.multi_scale,
return_targets=False)
# dataset.n = 4
# dataset.images = dataset.images[:dataset.n]
# dataset = SSDatasetYOLO( root_dir='../../../Data/SS/',
# class_file='../../../Data/SS/ss.names',
# anchors=model.anchors,
# dataset='test',
# skip_truncated=False,
# skip_difficult=False,
# mu=[0.141],
# sigma=[0.469],
# mode='stft',
# return_targets=True
# )
# dataset.images = [d for i, d in enumerate(dataset.images) if i in [2, 6, 9, 12, 14, 21, 23, 31]]
dataset.images = [d for i, d in enumerate(dataset.images) if i < 20]
dataset.n = len(dataset.images)
model = pickle.load(open('YOLOv2_120.pkl', 'rb'))
draw_vector_bboxes(model, dataset)
from torch import optim
from dataset import PascalDatasetYOLO
from layers import *
from darknet import YOLO
from utils import step_decay_scheduler, set_random_seed
if __name__ == '__main__':
set_random_seed(12345)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
train = False
predict = True
model = YOLO(name='YOLOv3-Tiny',
model='models/yolov3-tiny-voc.cfg',
device=device)
train_data = PascalDatasetYOLO(
root_dir=[
'../../../Data/VOCdevkit/VOC2007/',
'../../../Data/VOCdevkit/VOC2012/'
],
class_file='../../../Data/VOCdevkit/voc.names',
dataset=['trainval', 'trainval'],
batch_size=model.batch_size // model.subdivisions,
image_size=model.default_image_size,
anchors=model.anchors,
strides=model.strides,
do_transforms=True,
multi_scale=model.multi_scale)
import pickle
if __name__ == '__main__':
set_random_seed(12345)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
train = False
freeze = False
predict = True
model = YOLO(
name='YOLOv2',
# model='models/yolov2-tiny-voc.cfg',
model='models/yolov2-voc.cfg',
device=device)
# torchsummary.summary(model, (model.channels, *model.default_image_size), device=device)
train_data = PascalDatasetYOLO(
root_dir=[
'../../../Data/VOCdevkit/VOC2007/',
'../../../Data/VOCdevkit/VOC2012/'
],
class_file='../../../Data/VOCdevkit/voc.names',
dataset=['trainval', 'trainval'],
batch_size=model.batch_size // model.subdivisions,
image_size=model.default_image_size,
anchors=model.anchors,