def main(args):
cfg = setup(args)
model = DefaultTrainer.build_model(cfg)
modules = model.get_conv_bn_modules()
for m in modules:
print(m.weight.size())
load_darknet_weights(args.initial_weights, modules)
save_path = os.path.join(args.output_dir, "yolov3.pth")
torch.save(model.state_dict(), save_path)
print("model save to", save_path)
def __init__(self, config_file, mode='caffe', extract_mode='roi_feats', min_max_boxes='10,50'):
args = {}
args['config_file'] = config_file
args['mode'] = mode
args['extract_mode'] = extract_mode
args['min_max_boxes'] = min_max_boxes
args['eval_only'] = True
args = Pack(args)
cfg = setup(args)
self.model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(self.model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=False
)
self.model.eval()
self.cfg = cfg
self.vg_objects = load_vocab_file('objects_vocab.txt')
def evaluate(dataset):
register_one_set(dataset)
cfg = get_my_cfg()
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
model = DefaultTrainer.build_model(
cfg) # just built the model without weights
checkpoiner = DetectionCheckpointer(model, cfg.OUTPUT_DIR)
checkpoiner.resume_or_load(
cfg.MODEL.WEIGHTS, resume=False) # loaded the weights we had trained
evaluator = COCOEvaluator(dataset, ("bbox", ),
False,
output_dir=os.path.join("output", "evaluate"))
loader = build_detection_test_loader(cfg, dataset)
print(inference_on_dataset(model, loader, evaluator))
def main(args):
# if args.unitest:
# return unitest()
cfg = setup(args)
for d in ["train", 'val']:
# train for 6998images , val for 1199 images
DatasetCatalog.register("chefCap_" + d,
lambda d=d: get_chefcap_image_dicts())
MetadataCatalog.get("chefCap_" + d).set(
thing_classes=list(things_class_dict.keys()))
if d == 'val':
MetadataCatalog.get("chefCap_val").evaluator_type = "pascal_voc"
MetadataCatalog.get("chefCap_val").year = 2012
MetadataCatalog.get(
"chefCap_val"
).dirname = "/opt/work/chefCap/detectron2_fasterrcnn/data"
# for d in ["/opt/work/chefCap/data/ziped/Making-PascalVOC-export/"]:
# DatasetCatalog.register("chefCap_val",
# lambda d=d: get_chefcap_image_dicts(d))
# MetadataCatalog.get("chefCap_val").set(
# thing_classes=['face-head', 'mask-head', 'face-cap', 'mask-cap'])
# MetadataCatalog.get("chefCap_val").evaluator_type = "pascal_voc"
# MetadataCatalog.get("chefCap_val").dirname = "/opt/work/chefCap/data/ziped/Making-PascalVOC-export/"
# MetadataCatalog.get("chefCap_val").year = 2012
if args.eval_only:
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
res = DefaultTrainer.test(cfg, model)
if cfg.TEST.AUG.ENABLED:
res.update(DefaultTrainer.test_with_TTA(cfg, model))
if comm.is_main_process():
verify_results(cfg, res)
return res
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=args.resume)
# if cfg.TEST.AUG.ENABLED:
# trainer.register_hooks(
# [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
# )
return trainer.train()
def load_model_img():
global net_img
global cfg
time_start = time.time()
print('\nLoad image feature extraction model')
logging.disable()
parser = argparse.ArgumentParser()
parser.add_argument(
'--config-file',
default='model/extract-bua-caffe-r101.yaml',
metavar='FILE',
help='path to config file',
)
parser.add_argument('--mode',
default='caffe',
type=str,
help='bua_caffe, ...')
parser.add_argument(
'--resume',
action='store_true',
help='whether to attempt to resume from the checkpoint directory',
)
parser.add_argument(
'opts',
help='Modify config options using the command-line',
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args([])
cfg = setup(args)
MIN_BOXES = 10
MAX_BOXES = 100
CONF_THRESH = 0.2
net_img = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(net_img, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
net_img.eval()
time_end = time.time()
print('Finish load net model!')
print('Model load time: {:.3f}s\n'.format(time_end - time_start))
def __init__(self, min_boxes=3, max_boxes=10, threshold=0.5):
config_file = 'bottom-up-attention.pytorch/configs/bua-caffe/extract-bua-caffe-r101.yaml'
self._cfg = get_cfg()
add_bottom_up_attention_config(self._cfg, True)
self._cfg.merge_from_file(config_file)
self._cfg.MODEL.DEVICE = 'cpu'
self._model = DefaultTrainer.build_model(self._cfg)
DetectionCheckpointer(self._model, save_dir=self._cfg.OUTPUT_DIR).resume_or_load(self._cfg.MODEL.WEIGHTS)
self._model.eval()
self._min_boxes = min_boxes
self._max_boxes = max_boxes
self._threshold = threshold
self._classes = ['__background__']
with open(os.path.join('bottom-up-attention.pytorch', 'evaluation', 'objects_vocab.txt')) as f:
for object in f.readlines():
self._classes.append(object.split(',')[0].lower().strip())
def main(args):
cfg = setup(args)
show = True
register_openlogo(cfg.DATASETS.TRAIN[0], "datasets/data/openlogo",
"trainval", "supervised_imageset")
register_openlogo(cfg.DATASETS.TEST[0], "datasets/data/openlogo", "test",
"supervised_imageset")
trainer = DefaultTrainer(cfg)
evaluator = OpenLogoDetectionEvaluator(cfg.DATASETS.TEST[0])
if args.eval_only:
model = trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
if show:
visualize(cfg, amount=20)
res = trainer.test(cfg, model, evaluators=[evaluator])
if comm.is_main_process():
verify_results(cfg, res)
if cfg.TEST.AUG.ENABLED:
res.update(trainer.test_with_TTA(cfg, model))
return res
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=args.resume)
if cfg.TEST.AUG.ENABLED:
trainer.register_hooks([
hooks.EvalHook(0,
lambda: trainer.test_with_TTA(cfg, trainer.model))
])
return trainer.train()
def main(args):
cfg = setup(args)
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume
)
save_onnx_name = cfg.MODEL.WEIGHTS.replace('.pth', '.onnx')
model.eval()
inputs = torch.randn(1, 3, cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST).to(model.device)
with torch.no_grad():
torch.onnx.export(
model,
inputs,
save_onnx_name,
export_params=True,
verbose=False,
training=False,
do_constant_folding=True,
input_names=['input'],
# output_names=['output'],
# opset_version=9,
)
cfg.DATASETS.TEST = ("openimages_val", )
cfg.DATASETS.TRAIN = ("openimages_train", )
predictor = DefaultPredictor(cfg)
# trainer = DefaultTrainer(cfg)
if args.mode == 'infer':
val_dicts = get_val_dicts()
for d in random.sample(val_dicts, 10):
im = cv2.imread(d['file_name'])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=openimages_val_metadata,
scale=0.8,
instance_mode=ColorMode.IMAGE_BW)
v = v.draw_instance_predictions(outputs['instances'].to('cpu'))
cv2.imshow('image', v.get_image()[:, :, ::-1])
cv2.waitKey(0)
if args.mode == 'evaluate':
print(cfg.DATASETS.TEST[0])
trainer = DefaultTrainer(cfg)
model = trainer.build_model(cfg)
DetectionCheckpointer(model).load(args.model_pth)
evaluator = COCOEvaluator(cfg.DATASETS.TEST[0],
cfg,
False,
output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "openimages_val")
inference_on_dataset(model, val_loader, evaluator)
def build_tsv(gpu_id: int = 0):
args = Argument()
cfg = setup(args)
# os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model).resume_or_load(
cfg.MODEL.WEIGHTS,
resume=False,
)
model.eval()
# Set up the simulator
print("Starting sim")
sim = MatterSim.Simulator()
sim.setCameraResolution(WIDTH, HEIGHT)
sim.setCameraVFOV(math.radians(VFOV))
sim.setDiscretizedViewingAngles(False)
sim.setBatchSize(1)
sim.setPreloadingEnabled(False)
sim.initialize()
print("Sim has start")
# import pdb
# pdb.set_trace()
print("STARTING")
with open(OUTFILE, "wt") as tsvfile:
writer = csv.DictWriter(tsvfile, delimiter="\t", fieldnames=TSV_FIELDNAMES)
# Loop all the viewpoints in the simulator
print("loading viewpoints")
viewpoint_ids = load_viewpointids()
for scanId, viewpointId in tqdm(viewpoint_ids):
# Loop all discretized views from this location
ims = []
sim.newEpisode(
[scanId], [viewpointId], [0], [math.radians(ELEVATION_START)]
)
for ix in range(VIEWPOINT_SIZE):
state = sim.getState()[0]
# Transform and save generated image
ims.append(transform_img(state.rgb))
# Build state
if ix == 0:
record = {
"scanId": state.scanId,
"viewpointId": state.location.viewpointId,
"viewHeading": np.zeros(VIEWPOINT_SIZE, dtype=np.float32),
"viewElevation": np.zeros(VIEWPOINT_SIZE, dtype=np.float32),
"image_h": HEIGHT,
"image_w": WIDTH,
"vfov": VFOV,
}
record["viewHeading"][ix] = state.heading
record["viewElevation"][ix] = state.elevation
# Move the sim viewpoint so it ends in the same place
elev = 0.0
heading_chg = math.pi * 2 / VIEWS_PER_SWEEP
view = ix % VIEWS_PER_SWEEP
sweep = ix // VIEWS_PER_SWEEP
if view + 1 == VIEWS_PER_SWEEP: # Last viewpoint in sweep
elev = math.radians(ELEVATION_INC)
sim.makeAction([0], [heading_chg], [elev])
# Run detection
for ix in range(VIEWPOINT_SIZE):
get_detections_from_im(record, model, ims[ix], cfg)
filter(record, MAX_TOTAL_BOXES)
for k, v in record.items():
if isinstance(v, np.ndarray):
record[k] = str(base64.b64encode(v), "utf-8")
writer.writerow(record)
def extract_feat(split_idx, img_list, cfg, args, actor: ActorHandle):
num_images = len(img_list)
print('Number of images on split{}: {}.'.format(split_idx, num_images))
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
model.eval()
generate_npz_list = []
for im_file in (img_list):
if os.path.exists(
os.path.join(args.output_dir,
im_file.split('.')[0] + '.npz')):
actor.update.remote(1)
continue
im = cv2.imread(os.path.join(args.image_dir, im_file))
if im is None:
print(os.path.join(args.image_dir, im_file), "is illegal!")
actor.update.remote(1)
continue
dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
# extract roi features
if cfg.MODEL.BUA.EXTRACTOR.MODE == 1:
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
boxes = [box.tensor.cpu() for box in boxes]
scores = [score.cpu() for score in scores]
features_pooled = [feat.cpu() for feat in features_pooled]
if not attr_scores is None:
attr_scores = [attr_score.cpu() for attr_score in attr_scores]
generate_npz_list.append(
generate_npz.remote(1, actor, args, cfg, im_file, im,
dataset_dict, boxes, scores,
features_pooled, attr_scores))
# extract bbox only
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2:
with torch.set_grad_enabled(False):
boxes, scores = model([dataset_dict])
boxes = [box.cpu() for box in boxes]
scores = [score.cpu() for score in scores]
generate_npz_list.append(
generate_npz.remote(2, actor, args, cfg, im_file, im,
dataset_dict, boxes, scores))
# extract roi features by bbox
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3:
if not os.path.exists(
os.path.join(args.bbox_dir,
im_file.split('.')[0] + '.npz')):
actor.update.remote(1)
continue
bbox = torch.from_numpy(
np.load(
os.path.join(args.bbox_dir,
im_file.split('.')[0] +
'.npz'))['bbox']) * dataset_dict['im_scale']
proposals = Instances(dataset_dict['image'].shape[-2:])
proposals.proposal_boxes = BUABoxes(bbox)
dataset_dict['proposals'] = proposals
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
boxes = [box.tensor.cpu() for box in boxes]
scores = [score.cpu() for score in scores]
features_pooled = [feat.cpu() for feat in features_pooled]
if not attr_scores is None:
attr_scores = [
attr_score.data.cpu() for attr_score in attr_scores
]
generate_npz_list.append(
generate_npz.remote(3, actor, args, cfg, im_file, im,
dataset_dict, boxes, scores,
features_pooled, attr_scores))
ray.get(generate_npz_list)
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection2 Inference")
parser.add_argument(
"--config-file",
default="configs/bua-caffe/extract-bua-caffe-r101.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--mode",
default="caffe",
type=str,
help="bua_caffe, ...")
parser.add_argument('--out-dir',
dest='output_dir',
help='output directory for features',
default="features")
parser.add_argument('--image-dir',
dest='image_dir',
help='directory with images',
default="image")
parser.add_argument(
"--resume",
action="store_true",
help="whether to attempt to resume from the checkpoint directory",
)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg = setup(args)
MIN_BOXES = 10
MAX_BOXES = 100
CONF_THRESH = 0.2
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
# Extract features.
imglist = os.listdir(args.image_dir)
num_images = len(imglist)
print('Number of images: {}.'.format(num_images))
model.eval()
for im_file in tqdm.tqdm(imglist):
im = cv2.imread(os.path.join(args.image_dir, im_file))
dataset_dict = get_image_blob(im)
with torch.set_grad_enabled(False):
# boxes, scores, features_pooled = model([dataset_dict])
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
dets = boxes[0].tensor.cpu() / dataset_dict['im_scale']
scores = scores[0].cpu()
feats = features_pooled[0].cpu()
max_conf = torch.zeros((scores.shape[0])).to(scores.device)
for cls_ind in range(1, scores.shape[1]):
cls_scores = scores[:, cls_ind]
keep = nms(dets, cls_scores, 0.3)
max_conf[keep] = torch.where(cls_scores[keep] > max_conf[keep],
cls_scores[keep], max_conf[keep])
keep_boxes = torch.nonzero(max_conf >= CONF_THRESH).flatten()
if len(keep_boxes) < MIN_BOXES:
keep_boxes = torch.argsort(max_conf, descending=True)[:MIN_BOXES]
elif len(keep_boxes) > MAX_BOXES:
keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES]
image_feat = feats[keep_boxes]
image_bboxes = dets[keep_boxes]
image_objects_conf = np.max(scores[keep_boxes].numpy(), axis=1)
image_objects = np.argmax(scores[keep_boxes].numpy(), axis=1)
if cfg.MODEL.BUA.ATTRIBUTE_ON:
attr_scores = attr_scores[0].cpu()
image_attrs_conf = np.max(attr_scores[keep_boxes].numpy(), axis=1)
image_attrs = np.argmax(attr_scores[keep_boxes].numpy(), axis=1)
info = {
'image_id': im_file.split('.')[0],
'image_h': np.size(im, 0),
'image_w': np.size(im, 1),
'num_boxes': len(keep_boxes),
'objects_id': image_objects,
'objects_conf': image_objects_conf,
'attrs_id': image_attrs,
'attrs_conf': image_attrs_conf,
}
else:
info = {
'image_id': im_file.split('.')[0],
'image_h': np.size(im, 0),
'image_w': np.size(im, 1),
'num_boxes': len(keep_boxes),
'objects_id': image_objects,
'objects_conf': image_objects_conf
}
output_file = os.path.join(args.output_dir, im_file.split('.')[0])
np.savez_compressed(output_file,
x=image_feat,
bbox=image_bboxes,
num_bbox=len(keep_boxes),
image_h=np.size(im, 0),
image_w=np.size(im, 1),
info=info)
def extract_feat(split_idx, img_list, cfg, args, actor: ActorHandle):
num_images = len(img_list)
print('Number of images on split{}: {}.'.format(split_idx, num_images))
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
model.eval()
register()
print(cfg.MODEL.BUA.EXTRACTOR.MODE)
for im_file in (img_list):
# if os.path.exists(os.path.join(special_output_dir, im_file.split('.')[0]+'.npz')):
# actor.update.remote(1)
# continue
image_id = im_file.split('.')[0] # xxx.jpg
dump_folder = os.path.join(args.output_dir, str(image_id) + ".npz")
if os.path.exists(
os.path.join(args.output_dir,
im_file.split('.')[0] + '.npz')):
actor.update.remote(1)
continue
# else:
# start = True
# if not start:
# actor.update.remote(1)
# continue
im = cv2.imread(os.path.join(args.image_dir, im_file))
illegal = False
if im is None:
illegal = True
elif im.shape[-1] != 3:
illegal = True
elif max(im.shape[:2]) / min(im.shape[:2]) > 10 or max(
im.shape[:2]) < 25:
illegal = True
if illegal:
# print(os.path.join(args.image_dir, im_file), "is illegal!")
actor.update.remote(1)
continue
# dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
pixel_mean = cfg.MODEL.PIXEL_MEAN if args.mode == "caffe" else 0.0
image_h = np.size(im, 0)
image_w = np.size(im, 1)
dataset_dict = get_image_blob(im, pixel_mean)
# extract roi features
if cfg.MODEL.BUA.EXTRACTOR.MODE == 1:
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model_inference(
model, [dataset_dict], 1, dump_folder, image_h,
image_w, True)
else:
boxes, scores, features_pooled = model_inference(
model, [dataset_dict], 1, dump_folder, image_h,
image_w)
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2:
with torch.set_grad_enabled(False):
boxes, scores, _ = model_inference(model, [dataset_dict], 2,
dump_folder, image_h,
image_w)
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3: # extract roi features by bbox
npy = False
if os.path.exists(
os.path.join(args.bbox_dir,
im_file.split('.')[0] + '.npy')):
npy = True
elif not os.path.exists(
os.path.join(args.bbox_dir,
im_file.split('.')[0] + '.npz')):
actor.update.remote(1)
continue
if npy:
try:
bbox = torch.from_numpy(
np.load(os.path.join(args.bbox_dir,
im_file.split('.')[0] + '.npy'),
allow_pickle=True).tolist()
['bbox']) * dataset_dict['im_scale']
except Exception as e:
print(e)
continue
actor.update.remote(1)
else:
bbox = torch.from_numpy(
np.load(
os.path.join(
args.bbox_dir,
im_file.split('.')[0] +
'.npz'))['bbox']) * dataset_dict['im_scale']
proposals = Instances(dataset_dict['image'].shape[-2:])
proposals.proposal_boxes = BUABoxes(bbox)
dataset_dict['proposals'] = proposals
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model_inference(
model, [dataset_dict], 3, dump_folder, image_h,
image_w, True)
else:
boxes, scores, features_pooled = model_inference(
model, [dataset_dict], 3, dump_folder, image_h,
image_w)
else:
raise Exception("extract mode not supported.")
actor.update.remote(1)
# In[4]:
config_file = '../configs/bua-caffe/extract-bua-caffe-r101.yaml'
cfg = get_cfg()
add_bottom_up_attention_config(cfg, True)
cfg.merge_from_file(config_file)
cfg.freeze()
# In[5]:
MIN_BOXES = 10
MAX_BOXES = 20
CONF_THRESH = 0.4
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
'../' + cfg.MODEL.WEIGHTS, resume=True)
model.eval()
# In[6]:
from glob import glob
# im_files = ['../datasets/demo/004545.jpg']
im_files = sorted(glob("../../sample_data/*.png"))
for im_file in im_files:
im = cv2.imread(im_file)
dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection2 Inference")
parser.add_argument(
"--config-file",
default="configs/bua-caffe/extract-bua-caffe-r101.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--mode",
default="caffe",
type=str,
help="bua_caffe, ...")
parser.add_argument('--out-dir',
dest='output_dir',
help='output directory for features',
default="data/features")
parser.add_argument('--image-dir',
dest='image_dir',
help='directory with images',
default="data/coco_img")
parser.add_argument('--gt-bbox-dir',
dest='gt_bbox_dir',
help='directory with gt-bbox',
default="data/bbox")
parser.add_argument(
"--resume",
action="store_true",
help="whether to attempt to resume from the checkpoint directory",
)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg = setup(args)
MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
# Extract features.
imglist = os.listdir(args.image_dir)
num_images = len(imglist)
print('Number of images: {}.'.format(num_images))
model.eval()
for im_file in tqdm.tqdm(imglist):
if os.path.exists(
os.path.join(args.output_dir,
im_file.split('.')[0] + '.npz')):
continue
im = cv2.imread(os.path.join(args.image_dir, im_file))
if im is None:
print(os.path.join(args.image_dir, im_file), "is illegal!")
continue
dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
# extract roi features
if cfg.MODEL.BUA.EXTRACTOR.MODE == 1:
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
save_roi_features(args, cfg, im_file, im, dataset_dict, boxes,
scores, features_pooled, attr_scores)
# extract bbox only
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 2:
with torch.set_grad_enabled(False):
boxes, scores = model([dataset_dict])
save_bbox(args, cfg, im_file, im, dataset_dict, boxes, scores)
# extract roi features by gt bbox
elif cfg.MODEL.BUA.EXTRACTOR.MODE == 3:
if not os.path.exists(
os.path.join(args.gt_bbox_dir,
im_file.split('.')[0] + '.npz')):
continue
bbox = torch.from_numpy(
np.load(
os.path.join(args.gt_bbox_dir,
im_file.split('.')[0] +
'.npz'))['bbox']) * dataset_dict['im_scale']
proposals = Instances(dataset_dict['image'].shape[-2:])
proposals.proposal_boxes = BUABoxes(bbox)
dataset_dict['proposals'] = proposals
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
save_roi_features_by_gt_bbox(args, cfg, im_file, im, dataset_dict,
boxes, scores, features_pooled,
attr_scores)
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection2 Inference")
parser.add_argument(
"--config-file",
default="configs/bua-caffe/extract-bua-caffe-r101-fix36.yaml",
metavar="FILE",
help="path to config file",
)
# --image-dir or --image
parser.add_argument('--image-dir',
dest='image_dir',
help='directory with images',
default="datasets/demos")
parser.add_argument(
'--image', dest='image',
help='image') # e.g. datasets/demos/COCO_val2014_000000060623.jpg
parser.add_argument("--mode",
default="caffe",
type=str,
help="bua_caffe, ...")
parser.add_argument('--out-dir',
dest='output_dir',
help='output directory for features',
default="features")
parser.add_argument('--out-name',
dest='output_name',
help='output file name for features',
default="demos")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg = setup(args)
model = DefaultTrainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=True)
model.eval()
# Extract features.
if args.image:
imglist = [args.image]
else:
imglist = os.listdir(args.image_dir)
imglist = [os.path.join(args.image_dir, fn) for fn in imglist]
num_images = len(imglist)
print('Number of images: {}.'.format(num_images))
imglist.sort()
MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH
classes = []
with open(os.path.join('evaluation/objects_vocab.txt')) as f:
for object in f.readlines():
names = [n.lower().strip() for n in object.split(',')]
classes.append(names[0])
attributes = []
with open(os.path.join('evaluation/attributes_vocab.txt')) as f:
for att in f.readlines():
names = [n.lower().strip() for n in att.split(',')]
attributes.append(names[0])
classes = np.array(classes)
attributes = np.array(attributes)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
with h5py.File(os.path.join(args.output_dir, '%s_fc.h5' % args.output_name), 'a') as file_fc, \
h5py.File(os.path.join(args.output_dir, '%s_att.h5' % args.output_name), 'a') as file_att, \
h5py.File(os.path.join(args.output_dir, '%s_box.h5' % args.output_name), 'a') as file_box:
informations = {}
try:
for im_file in tqdm.tqdm(imglist):
img_nm = os.path.basename(im_file)
im = cv2.imread(im_file)
if im is None:
print(im_file, "is illegal!")
continue
dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
# extract roi features
attr_scores = None
with torch.set_grad_enabled(False):
if cfg.MODEL.BUA.ATTRIBUTE_ON:
boxes, scores, features_pooled, attr_scores = model(
[dataset_dict])
else:
boxes, scores, features_pooled = model([dataset_dict])
dets = boxes[0].tensor.cpu() / dataset_dict['im_scale']
scores = scores[0].cpu()
feats = features_pooled[0].cpu()
max_conf = torch.zeros((scores.shape[0])).to(scores.device)
for cls_ind in range(1, scores.shape[1]):
cls_scores = scores[:, cls_ind]
keep = nms(dets, cls_scores, 0.3)
max_conf[keep] = torch.where(
cls_scores[keep] > max_conf[keep], cls_scores[keep],
max_conf[keep])
keep_boxes = torch.nonzero(max_conf >= CONF_THRESH).flatten()
if len(keep_boxes) < MIN_BOXES:
keep_boxes = torch.argsort(max_conf,
descending=True)[:MIN_BOXES]
elif len(keep_boxes) > MAX_BOXES:
keep_boxes = torch.argsort(max_conf,
descending=True)[:MAX_BOXES]
image_feat = feats[keep_boxes].numpy()
image_bboxes = dets[keep_boxes].numpy()
image_objects_conf = np.max(scores[keep_boxes].numpy()[:, 1:],
axis=1)
image_objects = classes[np.argmax(
scores[keep_boxes].numpy()[:, 1:], axis=1)]
info = {
'image_name': img_nm,
'image_h': np.size(im, 0),
'image_w': np.size(im, 1),
'num_boxes': len(keep_boxes),
'objects': image_objects,
'objects_conf': image_objects_conf
}
if attr_scores is not None:
attr_scores = attr_scores[0].cpu()
image_attrs_conf = np.max(
attr_scores[keep_boxes].numpy()[:, 1:], axis=1)
image_attrs = attributes[np.argmax(
attr_scores[keep_boxes].numpy()[:, 1:], axis=1)]
info['attrs'] = image_attrs
info['attrs_conf'] = image_attrs_conf
file_fc.create_dataset(img_nm, data=image_feat.mean(0))
file_att.create_dataset(img_nm, data=image_feat)
file_box.create_dataset(img_nm, data=image_bboxes)
informations[img_nm] = info
finally:
file_fc.close()
file_att.close()
file_box.close()
pickle.dump(
informations,
open(
os.path.join(args.output_dir,
'%s_info.pkl' % args.output_name), 'wb'))
print(
'--------------------------------------------------------------------'
)
def test_equivalence(cfg, infile, outfile):
from deeplab3.modeling.backbone.resnet import ResNet101
from detectron2.checkpoint.detection_checkpoint import DetectionCheckpointer
from detectron2.engine import DefaultTrainer
import torch.nn as nn
#Deeplab
model = ResNet101(cfg.MODEL.DEEPLAB, BatchNorm=nn.BatchNorm2d)
model.eval()
if torch.cuda.is_available():
checkpoint = torch.load(infile, map_location=torch.device('cuda'))
else:
checkpoint = torch.load(infile, map_location=torch.device('cpu'))
state_dict = {
k.replace('backbone.', ''): v
for k, v in checkpoint['state_dict'].items() if 'backbone' in k
}
model.load_state_dict(state_dict)
#Detectron2
model2 = DefaultTrainer.build_model(cfg)
model2.eval()
DetectionCheckpointer(model2,
save_dir=outfile).resume_or_load(cfg.MODEL.WEIGHTS)
#Check that the layers are identical
convs = []
downsample = []
bn = []
downsample_bn = []
for name, param in model.state_dict().items():
param_dict = {'name': name, 'param': param, 'size': param.size()}
if 'conv' in name:
convs.append(param_dict)
elif 'downsample.0' in name:
downsample.append(param_dict)
elif 'bn' in name:
bn.append(param_dict)
elif 'downsample.1' in name:
downsample_bn.append(param_dict)
else:
print(name)
convs2 = []
shortcut = []
bn2 = []
downsample_bn2 = []
for name, param in model2.backbone.state_dict().items():
param_dict = {'name': name, 'param': param, 'size': param.size()}
if 'conv' in name:
convs.append(param_dict)
elif 'downsample.0' in name:
downsample.append(param_dict)
elif 'bn' in name:
bn.append(param_dict)
elif 'downsample.1' in name:
downsample_bn.append(param_dict)
else:
print(name)
for p1, p2 in zip(convs, convs2):
print(p1['name'], p2['name'])
if (p1['size'] != p2['size']):
print(p1['size'], p2['size'])
raise ValueError('Sizes not equal')
assert (torch.all(torch.eq(p1['param'], p2['param'])))
for p1, p2 in zip(shortcut, downsample):
print(p1['name'], p2['name'])
if (p1['size'] != p2['size']):
print(p1['size'], p2['size'])
raise ValueError('Sizes not equal')
assert (torch.all(torch.eq(p1['param'], p2['param'])))
for p1, p2 in zip(bn, bn2):
print(p1['name'], p2['name'])
if (p1['size'] != p2['size']):
print(p1['size'], p2['size'])
raise ValueError('Sizes not equal')
assert (torch.all(torch.eq(p1['param'], p2['param'])))
for p1, p2 in zip(downsample_bn2, downsample_bn):
print(p1['name'], p2['name'])
if (p1['size'] != p2['size']):
print(p1['size'], p2['size'])
raise ValueError('Sizes not equal')
assert (torch.all(torch.eq(p1['param'], p2['param'])))
#Check that output is identical
input = torch.ones((1, 4, 640, 640), dtype=torch.float32)
output = model(input)
output2 = model2.backbone(input)
assert (torch.all(torch.eq(output['res5'], output2['res5'])))
assert (torch.all(torch.eq(output['res2'], output2['res2'])))
print("Passed all tests")
