def __init__(self, root, image_sets=('train', ), transform=None,
target_transform=VOCAnnotationTransform(dict(zip(BCCD_CLASSES, range(len(BCCD_CLASSES))))),
dataset_name='BCCD'):
super().__init__(root, tuple(), transform, target_transform, dataset_name)
for name in image_sets:
rootpath = self.root
for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')):
self.ids.append((rootpath, line.strip()))
def detect():
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# load net
# cfg = config.voc_af
input_size = [416, 416]
num_classes = len(VOC_CLASSES)
testset = VOCDetection(
args.voc_root, [('2007', 'test')],
BaseTransform(input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)), VOCAnnotationTransform())
# build model
if args.version == 'yolo':
from model.yolo import YOLO
net = YOLO(device,
input_size=input_size,
num_classes=num_classes,
trainable=False)
print('Let us test yolo on the VOC0712 dataset ......')
else:
print('Unknown Version !!!')
exit()
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.eval()
print('Finished loading model!')
net = net.to(device)
# evaluation
detect_net(net,
device,
testset,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
thresh=args.visual_threshold)
def showData():
dataDir = r'/mnt/hdisk1/trainSet/pytorch/antifollow'
trainData = VOCDetection(root=dataDir,
image_sets=[('2012', 'trainval')],
transform=SSDAugmentationTest(),
target_transform=VOCAnnotationTransform(
None, True),
dataset_name='VOC0712')
epoch_size = len(trainData)
data_loader = torchData.DataLoader(trainData,
2,
num_workers=1,
shuffle=False,
collate_fn=getNumpy_collate,
pin_memory=True)
batch_iterator = iter(data_loader)
picCount = 0
while (True):
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
break
#show img
for i in range(len(images)):
img = images[i]
imWidth = img.shape[1]
imHeight = img.shape[0]
imgColor = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
for rect in targets[i]:
cv2.rectangle(
imgColor,
(int(rect[0] * imWidth), int(rect[1] * imHeight)),
(int(rect[2] * imWidth), int(rect[3] * imHeight)),
VOC_COLOR_ID_MAP[int(rect[4])])
picCount += 1
print(picCount)
cv2.imshow("input", imgColor)
cv2.waitKey()
if __name__ == '__main__':
# 加载网络
num_classes = len(labelmap) + 1
net = build_ssd('test', 300, num_classes)
net.load_state_dict(torch.load(args.trained_model))
if torch.cuda.is_available():
net = net.cuda()
cudnn.benchmark = True
net.eval()
# 记载数据
dataset = VOCDetection(args.voc_root, [('2007', 'test')],
BaseTransform(300, (104, 117, 123)),
VOCAnnotationTransform())
# 测试结果
num_images = len(dataset)
all_boxes = [[[] for _ in range(num_images)]
for _ in range(len(labelmap) + 1)] # 用于保存所有符合条件的结果
with torch.no_grad():
for i in range(num_images):
img, gt, h, w = dataset.pull_item(i)
img = img.unsqueeze(0)
if torch.cuda.is_available():
img = img.cuda()
detections = net(img) # 得到结果,shape[1,21,200,5]
for j in range(1, detections.shape[1]): # 循环计算每个类别
dets = detections[
# build model
if args.version == 'yolo':
from model.yolo import YOLO
net = YOLO(device,
input_size=input_size,
num_classes=num_classes,
trainable=False)
print('Let us test yolo on the VOC0712 dataset ......')
else:
print('Unknown Version !!!')
exit()
# load net
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(
args.voc_root, [('2007', set_type)],
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)), VOCAnnotationTransform())
net = net.to(device)
# evaluation
detect_net(args.save_folder,
net,
device,
dataset,
args.top_k,
thresh=args.confidence_threshold)