def handle_ratio_prediction(prediction):
hboxes = prediction.bbox.data.numpy()
rboxes = prediction.get_field("rboxes").data.numpy()
ratios = prediction.get_field("ratios").data.numpy()
scores = prediction.get_field("scores").data.numpy()
labels = prediction.get_field("labels").data.numpy()
h_idx = np.where(ratios > 0.8)[0]
h = hboxes[h_idx]
hboxes_vtx = np.vstack([
h[:, 0], h[:, 1], h[:, 2], h[:, 1], h[:, 2], h[:, 3], h[:, 0], h[:, 3]
]).transpose((1, 0))
rboxes[h_idx] = hboxes_vtx
keep = poly_nms(
np.hstack([rboxes, scores[:, np.newaxis]]).astype(np.double), 0.1)
rboxes = rboxes[keep].astype(np.int32)
scores = scores[keep]
labels = labels[keep]
if len(rboxes) > 0:
rboxes = np.vstack(rboxes)
return rboxes, scores, labels
else:
return None, None, None
def main():
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
parser.add_argument(
"--config-file",
default="/home/asd/Project/AirplaneDetection/Gliding-vertex-Trainer/gliding_vertex-master/configs/glide/dota.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument(
"--ckpt",
help="The path to the checkpoint for test, default is the latest checkpoint.",
default="/home/asd/Project/AirplaneDetection/Gliding-vertex-Trainer/exp_dota/0909/model_final.pth",
)
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl", init_method="env://"
)
synchronize()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
save_dir = ""
logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
logger.info("Using {} GPUs".format(num_gpus))
logger.info(cfg)
logger.info("Collecting env info (might take some time)")
logger.info("\n" + collect_env_info())
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt
_ = checkpointer.load(ckpt, use_latest=args.ckpt is None)
"""
# Initialize mixed-precision if necessary
use_mixed_precision = cfg.DTYPE == 'float16'
amp_handle = amp.init(enabled=use_mixed_precision, verbose=cfg.AMP_VERBOSE)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt
_ = checkpointer.load(ckpt, use_latest=args.ckpt is None)
iou_types = ("bbox",)
if cfg.MODEL.MASK_ON:
iou_types = iou_types + ("segm",)
if cfg.MODEL.KEYPOINT_ON:
iou_types = iou_types + ("keypoints",)
output_folders = [None] * len(cfg.DATASETS.TEST)
dataset_names = cfg.DATASETS.TEST
if cfg.OUTPUT_DIR:
for idx, dataset_name in enumerate(dataset_names):
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
mkdir(output_folder)
output_folders[idx] = output_folder
data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
inference(
model,
data_loader_val,
dataset_name=dataset_name,
iou_types=iou_types,
box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
device=cfg.MODEL.DEVICE,
expected_results=cfg.TEST.EXPECTED_RESULTS,
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
output_folder=output_folder,
)
synchronize()
"""
from maskrcnn_benchmark.data.transforms.build import build_transforms
from PIL import Image
import torchvision.transforms.functional as F
transform = build_transforms( cfg, is_train=False )
img_dir = "/home/asd/Mission/GaoFen/airplane_detection/data/train/train_val/val/images"
res_dir = "/home/asd/Mission/GaoFen/airplane_detection/data/train/train_val/val/res"
model.eval()
# imgs = os.listdir( img_dir )
import glob
imgs = glob.glob(img_dir+"/*.tif")
for img in imgs:
img_path = os.path.join( img_dir, img )
img_pil = Image.open( img_path )
# for i in range( 360 ):
original_img = img_pil
# original_img = F.rotate( img_pil, 45, expand=True )
origin_w, origin_h = original_img.size
img, target = transform( original_img, None )
print(img.shape)
img = img.view( (1, img.shape[0], img.shape[1], img.shape[2] ) )
h, w = img.shape[2:]
if h % 32 != 0:
new_h = ( h // 32 + 1 ) * 32
else:
new_h = h
if w % 32 != 0:
new_w = ( w // 32 + 1 ) * 32
else:
new_w = w
ratio_w = 1. * new_w / w
ratio_h = 1. * new_h / h
padded_img = torch.zeros( (1, 3, new_h, new_w)).float()
padded_img[:, :, :h, :w] = img
prediction = model( padded_img.cuda() )[0]
prediction = prediction.resize((origin_w * ratio_w, origin_h * ratio_h))
hboxes = prediction.bbox.cpu()
rboxes = prediction.get_field( "rboxes" ).cpu()
ratios = prediction.get_field( "ratios" ).cpu()
scores = prediction.get_field( "scores" ).cpu()
# labels = prediction.get_field( "labels" ).cpu()
for rbox, ratio, score in zip( rboxes, ratios, scores ):
print( rbox )
print( ratio, score )
h_idx = ratios > 0.8
# print(hboxes)
h = hboxes[h_idx]
hboxes_vtx = torch.stack( [h[:, 0], h[:, 1], h[:, 2], h[:, 1], h[:, 2], h[:, 3], h[:, 0], h[:, 3]] ).permute( 1, 0 )
rboxes[h_idx] = hboxes_vtx
# rboxes = rboxes.data.numpy().astype( np.int32 )
rboxes = rboxes.data.numpy()
keep = poly_nms( np.hstack( [rboxes, scores.cpu().data.numpy()[:, np.newaxis]] ).astype( np.double ), 0.1 )
rboxes = rboxes[keep].astype( np.int32 )
scores = scores[keep]
hboxes = hboxes[keep]
keep = np.where( scores > 0.6 )
rboxes = rboxes[keep]
scores = scores[keep].tolist()
hboxes = hboxes[keep]
# rboxes = list( map( minAreaRect, rboxes ) )
if len( rboxes ) > 0:
rboxes = np.vstack( rboxes )
else:
rboxes = np.array( rboxes )
# vis( img_info["file_name"], rboxes )
# img = cv2.imread( original_img )
img = np.array( original_img.convert( "RGB" ) )[:, :, ::-1].copy()
cv2.polylines( img, rboxes.reshape(-1, 4, 2).astype( np.int32 ), True, (0, 255, 255), thickness=2, lineType=cv2.LINE_AA )
filename = img_path.split( "/" )[-1]
cv2.imwrite( "{}/{}".format( res_dir, filename ), img )
def handle_keypoint_prediction(prediction, ratio1, ratio2, ratio3, ratio4):
hboxes = prediction.bbox.data.numpy()
keypoints = prediction.get_field( "points" ).data.numpy()
# ratios = prediction.get_field( "ratios" ).data.numpy()
scores = prediction.get_field( "scores" ).data.numpy()
labels = prediction.get_field( "labels" ).data.numpy()
logits = prediction.get_field( "logits" ).data.numpy()
keypoints = np.array(keypoints)
keypoints = keypoints.reshape((-1, 8, 2))
logits = np.array(logits)
logits = logits.reshape((-1, 8))
logits = np.max(logits, axis=1)
rboxes = np.zeros((hboxes.shape[0], 8))
for i, keypoint in enumerate(keypoints):
# if labels[i] == 10 or labels[i] == 12:
# xmin = hboxes[i][0]
# ymin = hboxes[i][1]
# xmax = hboxes[i][2]
# ymax = hboxes[i][3]
# rboxes[i] = np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
# else:
rect = cv2.minAreaRect(np.float32(keypoint))
rect = cv2.boxPoints(rect).reshape(8)
poly = shgeo.Polygon(rect.reshape(4, 2))
rect1 = cv2.minAreaRect(np.float32(keypoint[0: 4]))
rect1 = cv2.boxPoints(rect1).reshape(8)
poly1 = shgeo.Polygon(rect1.reshape(4, 2))
poly1 = poly1.buffer(0.01)
poly1_1 = poly1.intersection(poly)
rect2 = np.float32(keypoint[0: 4]).reshape(8)
poly2 = shgeo.Polygon(rect2.reshape(4, 2))
poly2 = poly2.buffer(0.01)
poly2_1 = poly2.intersection(poly)
h_min = np.min(hboxes[i])
h_max = np.max(hboxes[i])
# scores[i] *= logits[i]
if h_min <=1 or h_max >=1024:
scores[i] *= 0.9
if poly2_1.area / (poly.area + poly2.area - poly2_1.area + 1e-10) < 0.30:
scores[i] = 0
ratio1 +=1
if poly2_1.area / (poly.area + poly2.area - poly2_1.area + 1e-10) > 0.75:
rboxes[i] = rect2
ratio2 += 1
elif poly1_1.area / (poly.area + poly1.area - poly1_1.area + 1e-10) > 0.75:
rboxes[i] = rect1
ratio3 += 1
else:
rboxes[i] = rect
scores[i] *= 0.9
ratio4 += 1
# rboxes[i] = postprocess(keypoint)
# for j, keypoint in enumerate(keypoints):
# weight = 0
# logit = logits[j]
# lines = []
# indexes = [[0, 5, 1],
# [1, 6, 2],
# [2, 7, 3],
# [3, 8, 0]]
# for index in indexes:
# fit = []
# for i in range(len(index)):
# if logit[index[i]] >= -4:
# fit.append(keypoint[index[i]])
# if len(fit) >= 2:
# fit = np.array(fit)
# output = cv2.fitLine(fit, cv2.DIST_L2, 0, 0.01, 0.01)
# k = output[1] / output[0]
# b = output[3] - k * output[2]
# x1 = 1
# y1 = k * x1 + b
# x2 = 100
# y2 = k * x2 + b
# lines.append(np.array([x1, y1, x2, y2]))
# else:
# weight = 1
# break
# if weight == 0:
# cross_point = []
# cross_point.append(getCrossPoint(lines[0], lines[1]))
# cross_point.append(getCrossPoint(lines[1], lines[2]))
# cross_point.append(getCrossPoint(lines[2], lines[3]))
# cross_point.append(getCrossPoint(lines[3], lines[0]))
# cross_point = np.array(cross_point)
# if -100 in cross_point:
# rect = cv2.minAreaRect(keypoint)
# rect = cv2.boxPoints(rect).reshape(8)
# else:
# rect = cross_point.reshape(8)
# # for keypoint in keypoints:
# else:
# rect = cv2.minAreaRect(keypoint)
# rect = cv2.boxPoints(rect).reshape(8)
# rboxes[j] = rect
# h_idx = np.where(ratios > 0.8)[0]
# h = hboxes[h_idx]
# hboxes_vtx = np.vstack( [h[:, 0], h[:, 1], h[:, 2], h[:, 1], h[:, 2], h[:, 3], h[:, 0], h[:, 3]] ).transpose((1,0))
# rboxes[h_idx] = hboxes_vtx
keep = poly_nms( np.hstack( [rboxes, scores[:, np.newaxis]] ).astype( np.double ), 0.1 )
rboxes = rboxes[keep].astype( np.int32 )
scores = scores[keep]
labels = labels[keep]
if len( rboxes ) > 0:
rboxes = np.vstack( rboxes )
return rboxes, scores, labels, ratio1, ratio2, ratio3, ratio4
else:
return None, None, None, ratio1, ratio2, ratio3, ratio4