337 lines
12 KiB
Python
337 lines
12 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
|
# pylint: disable=no-member
|
|
"""
|
|
TridentNet Training Script.
|
|
adf
|
|
This script is a simplified version of the training script in detectron2/tools.
|
|
"""
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import torch
|
|
# import tqdm
|
|
import cv2
|
|
import numpy as np
|
|
sys.path.append('detectron2')
|
|
|
|
import detectron2.utils.comm as comm
|
|
from detectron2.checkpoint import DetectionCheckpointer
|
|
from detectron2.data import build_detection_test_loader, build_detection_train_loader
|
|
from detectron2.config import get_cfg
|
|
from detectron2.engine import DefaultTrainer, default_setup, launch
|
|
from detectron2.evaluation import COCOEvaluator, verify_results
|
|
from detectron2.structures import Instances
|
|
from detectron2.modeling import build_model
|
|
|
|
from utils.utils import mkdir, save_features
|
|
from utils.extract_utils import get_image_blob, save_bbox, save_roi_features_by_bbox, save_roi_features, image_features
|
|
from utils.progress_bar import ProgressBar
|
|
from models import add_config
|
|
from models.bua.box_regression import BUABoxes
|
|
|
|
import ray
|
|
from ray.actor import ActorHandle
|
|
|
|
def switch_extract_mode(mode):
|
|
if mode == 'roi_feats':
|
|
switch_cmd = ['MODEL.BUA.EXTRACTOR.MODE', 1]
|
|
elif mode == 'bboxes':
|
|
switch_cmd = ['MODEL.BUA.EXTRACTOR.MODE', 2]
|
|
elif mode == 'bbox_feats':
|
|
switch_cmd = ['MODEL.BUA.EXTRACTOR.MODE', 3, 'MODEL.PROPOSAL_GENERATOR.NAME', 'PrecomputedProposals']
|
|
else:
|
|
print('Wrong extract mode! ')
|
|
exit()
|
|
return switch_cmd
|
|
|
|
def set_min_max_boxes(min_max_boxes):
|
|
if min_max_boxes == 'min_max_default':
|
|
return []
|
|
try:
|
|
min_boxes = int(min_max_boxes.split(',')[0])
|
|
max_boxes = int(min_max_boxes.split(',')[1])
|
|
except:
|
|
print('Illegal min-max boxes setting, using config default. ')
|
|
return []
|
|
cmd = ['MODEL.BUA.EXTRACTOR.MIN_BOXES', min_boxes,
|
|
'MODEL.BUA.EXTRACTOR.MAX_BOXES', max_boxes]
|
|
return cmd
|
|
|
|
def setup(args):
|
|
"""
|
|
Create configs and perform basic setups.
|
|
"""
|
|
cfg = get_cfg()
|
|
add_config(args, cfg)
|
|
cfg.merge_from_file(args.config_file)
|
|
cfg.merge_from_list(args.opts)
|
|
cfg.merge_from_list(switch_extract_mode(args.extract_mode))
|
|
cfg.merge_from_list(set_min_max_boxes(args.min_max_boxes))
|
|
cfg.freeze()
|
|
default_setup(cfg, args)
|
|
return cfg
|
|
|
|
def generate_npz(extract_mode, *args):
|
|
if extract_mode == 1:
|
|
save_roi_features(*args)
|
|
elif extract_mode == 2:
|
|
save_bbox(*args)
|
|
elif extract_mode == 3:
|
|
save_roi_features_by_bbox(*args)
|
|
else:
|
|
print('Invalid Extract Mode! ')
|
|
|
|
@ray.remote(num_gpus=1)
|
|
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))
|
|
#print(cfg)
|
|
model = DefaultTrainer.build_model(cfg)
|
|
print("111111111111111111111111111111111111111111111")
|
|
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
|
|
cfg.MODEL.WEIGHTS, resume=args.resume
|
|
)
|
|
print("222222222222222222222222222222222222222222222222")
|
|
model.eval()
|
|
|
|
for im_file in (img_list):
|
|
print("33333333333333333333333333333333333333333333")
|
|
# if os.path.exists(os.path.join(args.output_dir, im_file.split('.')[0]+'.npz')):
|
|
# actor.update.remote(1)
|
|
# continue
|
|
im = cv2.imread(im_file)
|
|
|
|
if im is None:
|
|
print(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
|
|
print("444444444444444444444444444444444")
|
|
with torch.set_grad_enabled(False):
|
|
if cfg.MODEL.BUA.ATTRIBUTE_ON:
|
|
print("55555555555555555555555555555")
|
|
boxes, scores, features_pooled, attr_scores = model([dataset_dict])
|
|
print("6666666666666666666666666")
|
|
else:
|
|
boxes, scores, features_pooled = model([dataset_dict])
|
|
boxes = [box.tensor.cpu() for box in boxes]
|
|
scores = [score.cpu() for score in scores]
|
|
#torch.set_printoptions(subprocess=True)
|
|
#print(features_pooled)
|
|
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]
|
|
print("77777777777777777777777777")
|
|
im_feat = image_features(cfg, im_file, im, dataset_dict,
|
|
boxes, scores, features_pooled, attr_scores)
|
|
print("888888888888888888888888")
|
|
actor.update.remote(1)
|
|
return im_feat
|
|
|
|
|
|
def feature(image_path):
|
|
parser = argparse.ArgumentParser(description="PyTorch Object Detection2 Inference")
|
|
parser.add_argument(
|
|
"--config-file",
|
|
default="configs/bua-caffe/extract-bua-caffe-r152.yaml",
|
|
metavar="FILE",
|
|
help="path to config file",
|
|
)
|
|
|
|
parser.add_argument('--num-cpus', default=1, type=int,
|
|
help='number of cpus to use for ray, 0 means no limit')
|
|
|
|
parser.add_argument('--gpus', dest='gpu_id', help='GPU id(s) to use',
|
|
default='0', type=str)
|
|
|
|
parser.add_argument("--mode", default="caffe", type=str, help="bua_caffe, ...")
|
|
|
|
parser.add_argument('--extract-mode', default='roi_feats', type=str,
|
|
help="'roi_feats', 'bboxes' and 'bbox_feats' indicates \
|
|
'extract roi features directly', 'extract bboxes only' and \
|
|
'extract roi features with pre-computed bboxes' respectively")
|
|
|
|
parser.add_argument('--min-max-boxes', default='min_max_default', type=str,
|
|
help='the number of min-max boxes of extractor')
|
|
|
|
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('--bbox-dir', dest='bbox_dir',
|
|
help='directory with bbox',
|
|
default="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)
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
|
|
num_gpus = len(args.gpu_id.split(','))
|
|
|
|
MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
|
|
MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
|
|
CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH
|
|
|
|
|
|
|
|
print("---------------------------------------------------"+str(args.num_cpus))
|
|
if args.num_cpus != 0:
|
|
#print(ray.init(num_cpus=args.num_cpus))
|
|
ray.init(num_cpus=args.num_cpus)
|
|
else:
|
|
#print(ray.init())
|
|
ray.init()
|
|
|
|
# Extract features.
|
|
imglist = []
|
|
imglist.append(image_path)
|
|
num_images = len(imglist)
|
|
print('Number of images: {}.'.format(num_images))
|
|
img_lists = [imglist[i::num_gpus] for i in range(num_gpus)]
|
|
|
|
pb = ProgressBar(len(imglist))
|
|
actor = pb.actor
|
|
|
|
print('Number of GPUs: {}.'.format(num_gpus))
|
|
extract_feat_list = []
|
|
#model = DefaultTrainer.build_model(cfg)
|
|
for i in range(num_gpus):
|
|
extract_feat_list.append(extract_feat.remote(i, img_lists[i], cfg, args, actor))
|
|
|
|
pb.print_until_done()
|
|
img_fe = ray.get(extract_feat_list)
|
|
#img_fe = np.array(img_fe)
|
|
#print(img_fe.shape)
|
|
ray.get(actor.get_counter.remote())
|
|
return img_fe[0]
|
|
|
|
def loadModel():
|
|
parser = argparse.ArgumentParser(description="PyTorch Object Detection2 Inference")
|
|
parser.add_argument(
|
|
"--config-file",
|
|
default="configs/bua-caffe/extract-bua-caffe-r152.yaml",
|
|
metavar="FILE",
|
|
help="path to config file",
|
|
)
|
|
|
|
parser.add_argument('--num-cpus', default=1, type=int,
|
|
help='number of cpus to use for ray, 0 means no limit')
|
|
|
|
parser.add_argument('--gpus', dest='gpu_id', help='GPU id(s) to use',
|
|
default='0', type=str)
|
|
|
|
parser.add_argument("--mode", default="caffe", type=str, help="bua_caffe, ...")
|
|
|
|
parser.add_argument('--extract-mode', default='roi_feats', type=str,
|
|
help="'roi_feats', 'bboxes' and 'bbox_feats' indicates \
|
|
'extract roi features directly', 'extract bboxes only' and \
|
|
'extract roi features with pre-computed bboxes' respectively")
|
|
|
|
parser.add_argument('--min-max-boxes', default='min_max_default', type=str,
|
|
help='the number of min-max boxes of extractor')
|
|
|
|
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('--bbox-dir', dest='bbox_dir',
|
|
help='directory with bbox',
|
|
default="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)
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
|
|
num_gpus = len(args.gpu_id.split(','))
|
|
|
|
MIN_BOXES = cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
|
|
MAX_BOXES = cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
|
|
CONF_THRESH = cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH
|
|
|
|
print("---------------------------------------------------" + str(args.num_cpus))
|
|
if args.num_cpus != 0:
|
|
# print(ray.init(num_cpus=args.num_cpus))
|
|
ray.init(num_cpus=args.num_cpus)
|
|
else:
|
|
# print(ray.init())
|
|
ray.init()
|
|
|
|
model = DefaultTrainer.build_model(cfg)
|
|
print("111111111111111111111111111111111111111111111")
|
|
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
|
|
cfg.MODEL.WEIGHTS, resume=args.resume
|
|
)
|
|
print("222222222222222222222222222222222222222222222222")
|
|
model.eval()
|
|
return model, cfg
|
|
|
|
def imgFeature(model, cfg, im_file):
|
|
print("33333333333333333333333333333333333333333333")
|
|
# if os.path.exists(os.path.join(args.output_dir, im_file.split('.')[0]+'.npz')):
|
|
# actor.update.remote(1)
|
|
# continue
|
|
im = cv2.imread(im_file)
|
|
if im is None:
|
|
print(im_file, "is illegal!")
|
|
return 0
|
|
dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)
|
|
# extract roi features
|
|
if cfg.MODEL.BUA.EXTRACTOR.MODE == 1:
|
|
attr_scores = None
|
|
print("444444444444444444444444444444444")
|
|
with torch.set_grad_enabled(False):
|
|
if cfg.MODEL.BUA.ATTRIBUTE_ON:
|
|
print("55555555555555555555555555555")
|
|
boxes, scores, features_pooled, attr_scores = model([dataset_dict])
|
|
print("6666666666666666666666666")
|
|
else:
|
|
boxes, scores, features_pooled = model([dataset_dict])
|
|
boxes = [box.tensor.cpu() for box in boxes]
|
|
scores = [score.cpu() for score in scores]
|
|
# torch.set_printoptions(subprocess=True)
|
|
# print(features_pooled)
|
|
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]
|
|
print("77777777777777777777777777")
|
|
im_feat = image_features(cfg, im_file, im, dataset_dict,
|
|
boxes, scores, features_pooled, attr_scores)
|
|
print("888888888888888888888888")
|
|
return im_feat
|
|
|
|
if __name__ == "__main__":
|
|
feature("./image/134206.jpg")
|