def __init__(self, imdb_name, class_name, datadir, transform, dst_size):
meta = imdb_name.split('_')
self._image_set = meta[1]
self._name = imdb_name
self._data_path = os.path.join(datadir, cfg.args.datadirname,
class_name)
self._data_dir = self._data_path
self._classes = cfg.label_names
self._num_classes = len(self._classes)
self._class_to_ind = dict(
list(zip(self._classes, list(range(self._num_classes)))))
self._image_ext = '.jpg'
self._salt = str(uuid.uuid4())
self._comp_id = 'comp4'
self.config = {'cleanup': True, 'use_salt': True}
self.transform = transform
self._dst_size = dst_size
self._image_indexes = []
self._image_names = []
self._annotations = []
self._vertices = yolo_utils.threed_vertices(class_name)
self._objpoints3d = yolo_utils.threed_corners(class_name)
if cfg.args.dataset_name == 'LINEMOD':
with open("cfgs/diameter_linemod.yml", 'r') as stream:
diam_data = yaml.load(stream)
self._objdiam = diam_data[cfg.args.class_name]
self.load_dataset()
def dataset_eval(detpath,
annopath,
imagesetfile,
classname,
cachedir):
# first load gt
if not os.path.isdir(cachedir):
os.makedirs(cachedir)
cachefile = os.path.join(cachedir, 'annots.pkl')
# read list of images
with open(imagesetfile, 'r') as f:
lines = f.readlines()
imagenames = [x.strip() for x in lines]
if not os.path.isfile(cachefile):
# load annots
recs = {}
for i, imagename in enumerate(imagenames):
recs[imagename] = parse_rec(annopath.format(imagename), os.path.join(cfg.DATA_DIR, cfg.args.datadirname, classname, 'PoseAnnotations', '{:s}.txt').format(imagename))
if i % 100 == 0:
print('Reading annotation for {:d}/{:d}'.format(
i + 1, len(imagenames)))
# save
print('Saving cached annotations to {:s}'.format(cachefile))
with open(cachefile, 'wb') as f:
pickle.dump(recs, f)
else:
# load
with open(cachefile, 'rb') as f:
recs = pickle.load(f)
if cfg.args.num_detection_points == 9:
objpoints3D = threed_corners(classname)
else:
objpoints3D = threed_correspondences(classname)
vertices = threed_vertices(classname, eval=True)
if cfg.args.dataset_name == 'LINEMOD':
with open("cfgs/diameter_linemod.yml", 'r') as stream:
diam_data = yaml.load(stream)
if cfg.args.bias_correction == 'mean':
with open("cfgs/bias_linemod_mean.yml", 'r') as stream:
bias_data = yaml.load(stream)
if cfg.args.bias_correction == 'median':
with open("cfgs/bias_linemod_median.yml", 'r') as stream:
bias_data = yaml.load(stream)
if cfg.args.bias_correction == 'mode':
with open("cfgs/bias_linemod_mode.yml", 'r') as stream:
bias_data = yaml.load(stream)
obj_diameter = diam_data[classname]
if cfg.args.bias_correction is not None:
obj_bias = bias_data[classname]
else:
obj_bias = None
# extract gt objects for this class
class_recs = {}
npos = 0
for imagename in imagenames:
R = [obj for obj in recs[imagename] if obj['name'] == _class_to_ind[classname]]
bbox = np.array([x['bbox'] for x in R])
rt = np.array([x['rt'] for x in R])
# det = [False] * len(R)
npos = npos + len(R)
class_recs[imagename] = {'bbox': bbox, 'rt': rt}
# read dets
detfile = detpath.format(classname)
with open(detfile, 'r') as f:
lines = f.readlines()
if any(lines) == 1:
splitlines = [x.strip().split(' ') for x in lines]
image_ids = [x[0] for x in splitlines]
# confidence = np.array([float(x[1]) for x in splitlines])
BB = np.array([[float(z) for z in x[2:]] for x in splitlines])
# go down dets and mark TPs and FPs
nd = len(image_ids)
tp_2d = {}
for thresh in px_threshold:
tp_2d[thresh] = {}
rec2d = {}
tp_3d = {}
twod_dists = []
threed_dists = []
for d in xrange(nd):
R = class_recs[image_ids[d]]
bb = BB[d, :].astype(float)
BBGT = R['bbox'].astype(float)
RTGT = R['rt'].astype(float)
if BBGT.size > 0:
pixel_dist, vertex_error, vertex_dist = reprojection_errors(BBGT, RTGT, bb, image_ids[d], objpoints3D,
vertices, classname, obj_diameter, obj_bias)
twod_dists.append(pixel_dist)
threed_dists.append(vertex_dist)
if not vertex_error:
tp_3d[image_ids[d]] = 1
# Estimating 2d projection error based on different thresholds
for thresh in px_threshold:
if pixel_dist < thresh:
tp_2d[thresh][image_ids[d]] = 1
# Moving detection images in correct and incorrect directories
if cfg.args.vis:
results_cache = cfg.test_output_dir + "/results2d_cache.txt"
fo = open(results_cache, "w")
for k, v in tp_2d[5].items():
fo.write(str(k) + ' '+ str(v) + '\n')
os.rename(cfg.test_output_dir + '/' + str(k) + '.jpg', cfg.test_output_dir + '/correct/' + str(k) + '.jpg')
results_cache = cfg.test_output_dir + "/results3d_cache.txt"
fo = open(results_cache, "w")
for k, v in tp_3d.items():
fo.write(str(k) + ' '+ str(v) + '\n')
# Estimating 2d projection error based on different thresholds
for thresh in px_threshold:
tp_2d[thresh] = sum(tp_2d[thresh].values())
tp_3d = sum(tp_3d.values())
for thresh in px_threshold:
rec2d[thresh] = tp_2d[thresh] / float(npos)
rec3d = tp_3d / float(npos)
else:
rec2d = -1
rec3d = -1
if cfg.args.log_3derror:
with open('3derror_logs/' + classname + '_3derror_logs.yml', 'w') as outfile:
yaml.dump(log_3derror, outfile, default_flow_style=False)
return rec2d, rec3d, twod_dists, threed_dists
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from gn_utils import resnet_concat_gn, resnet_gn
import cfgs.config_yolo6d as cfg
from utils.yolo6d import threed_vertices, threed_correspondences, threed_corners
import utils.network as net_utils
vertices_old = threed_vertices(cfg.label_names[0])
corners3d = threed_corners(cfg.label_names[0])
vertices = np.c_[np.array(vertices_old),
np.ones((len(vertices_old), 1))].transpose()
if cfg.args.num_detection_points == 9:
objpoints3D = threed_corners(cfg.args.class_name)
else:
objpoints3D = threed_correspondences(cfg.args.class_name)
class Cullnet(nn.Module):
def __init__(self):
super(Cullnet, self).__init__()
if cfg.args.cullnet_type == 'resnet50_gn':
model = resnet_gn.resnet50(pretrained=False)
self.vismodel = nn.Sequential(*list(model.children())[:-1])
num_final_in = model.fc.in_features
# linear
out_channels = 1
self.regression = net_utils.FC(num_final_in,
out_channels,
relu=False)
dataloader = DataLoader(imdb,
batch_size=args.batch_size,
shuffle=False,
**kwargs)
net1 = Darknet()
net2 = Cullnet()
net1.load_state_dict(checkpoint1['state_dict'])
net1.cuda()
net2.load_state_dict(checkpoint2['state_dict'])
net2.cuda()
print('load net succ...')
if cfg.args.num_detection_points == 9:
objpoints3D = yolo_utils.threed_corners(args.class_name)
else:
objpoints3D = yolo_utils.threed_correspondences(args.class_name)
corners_3d = yolo_utils.threed_corners(args.class_name)
vertices = yolo_utils.threed_vertices(args.class_name)
if cfg.args.dataset_name == 'LINEMOD':
with open("cfgs/diameter_linemod.yml", 'r') as stream:
diam_data = yaml.load(stream)
obj_diameter = diam_data[args.class_name]
imdb_cullnettest = CullNetDataset_class(cfg.imdb_test, args.class_name,
cfg.DATA_DIR,
yolo_utils.preprocess_test_cullnet,