def anchor_iou_target_opr(self, boxes, im_info, all_anchors,
rpn_bbox_offsets):
n = rpn_bbox_offsets.shape[0]
res = []
for i in range(n):
gtboxes = boxes[i, :im_info[i, 5].astype(np.int32)]
offsets = rpn_bbox_offsets[i].reshape(-1, 4).detach()
m = offsets.shape[0]
an, ac = all_anchors.shape[0], all_anchors.shape[1]
anchors = F.broadcast_to(F.expand_dims(all_anchors, 1),
(an, 2, ac)).reshape(-1, ac)
dtboxes = bbox_transform_inv_opr(anchors[:, :4], offsets[:, :4])
overlaps = box_overlap_opr(dtboxes, gtboxes[:, :4])
ignore_mask = 1 - F.equal(
gtboxes[:, 4], config.anchor_ignore_label).astype(np.float32)
ignore_mask = F.expand_dims(ignore_mask, axis=0)
overlaps = overlaps * ignore_mask
overlaps = overlaps.reshape(-1, 2,
overlaps.shape[1]).transpose(1, 0, 2)
a, b = overlaps[0], overlaps[1]
index = F.argmax(a, axis=1)
a = F.nn.indexing_one_hot(a, index, 1)
b = F.scatter(b, 1, index.reshape(-1, 1), F.zeros([b.shape[0], 1]))
index = F.argmax(b, axis=1)
b = F.nn.indexing_one_hot(b, index, 1)
value = F.expand_dims(F.stack([a, b], axis=1), axis=0)
res.append(value)
result = F.concat(res, 0)
return result
def fpn_anchor_target_opr_core_impl(gt_boxes,
im_info,
anchors,
allow_low_quality_matches=True):
ignore_label = config.ignore_label
# get the gt boxes
gtboxes = gt_boxes[:im_info[5].astype(np.int32)]
ignore_mask = F.equal(gtboxes[:, 4], config.ignore_label)
# find the valid gtboxes
_, index = F.cond_take(1 - ignore_mask > 0, ignore_mask)
valid_gt_boxes = gtboxes[index.astype(np.int32)]
# compute the iou matrix
overlaps = box_overlap_opr(anchors, valid_gt_boxes[:, :4])
# match the dtboxes
a_shp0 = anchors.shape[0]
argmax_overlaps = F.argmax(overlaps, axis=1)
max_overlaps = F.nn.indexing_one_hot(overlaps,
argmax_overlaps.astype(np.int32), 1)
labels = F.ones(a_shp0).astype(np.int32) * ignore_label
# set negative ones
labels = labels * (max_overlaps >= config.rpn_negative_overlap).astype(
np.float32)
# set positive ones
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
const_one = mge.tensor(1.0)
if allow_low_quality_matches:
# match the max gt
gt_max_overlaps = F.max(overlaps, axis=0)
gt_argmax_overlaps = F.argmax(overlaps, axis=0)
gt_argmax_overlaps = gt_argmax_overlaps.astype(np.int32)
max_overlaps[gt_argmax_overlaps] = 1.
m = gt_max_overlaps.shape[0]
argmax_overlaps[gt_argmax_overlaps] = F.linspace(0, m - 1,
m).astype(np.int32)
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
labels[fg_mask] = 1
# compute the bbox targets
bbox_targets = bbox_transform_opr(anchors,
valid_gt_boxes[argmax_overlaps, :4])
if config.rpn_bbox_normalize_targets:
std_opr = mge.tensor(config.bbox_normalize_stds[None, :]).to(
anchors.device)
mean_opr = mge.tensor(config.bbox_normalize_means[None, :]).to(
anchors.device)
minus_opr = mean_opr / std_opr
bbox_targets = bbox_targets / std_opr - minus_opr
return labels, bbox_targets
def test_neg_axis():
x = tensor(np.random.normal(0, 1, (32, 5)))
y = F.argmax(x, axis=-1)
yy = F.argmax(x, axis=1)
np.testing.assert_equal(y.numpy(), yy.numpy())
y = F.argmax(x, axis=(-1, -2))
yy = F.argmax(x, axis=(0, 1))
np.testing.assert_equal(y.numpy(), yy.numpy())
y = F.argmin(x, axis=(-1, -2))
yy = F.argmin(x, axis=(0, 1))
np.testing.assert_equal(y.numpy(), yy.numpy())
def get_ground_truth(self, anchors, batched_gt_boxes, batched_valid_gt_box_number):
total_anchors = anchors.shape[0]
labels_cat_list = []
bbox_targets_list = []
for b_id in range(self.batch_size):
gt_boxes = batched_gt_boxes[b_id, : batched_valid_gt_box_number[b_id]]
overlaps = layers.get_iou(anchors, gt_boxes[:, :4])
argmax_overlaps = F.argmax(overlaps, axis=1)
max_overlaps = overlaps.ai[
F.linspace(0, total_anchors - 1, total_anchors).astype(np.int32),
argmax_overlaps,
]
labels = mge.tensor([-1]).broadcast(total_anchors)
labels = labels * (max_overlaps >= self.cfg.negative_thresh)
labels = labels * (max_overlaps < self.cfg.positive_thresh) + (
max_overlaps >= self.cfg.positive_thresh
)
bbox_targets = self.box_coder.encode(
anchors, gt_boxes.ai[argmax_overlaps, :4]
)
labels_cat = gt_boxes.ai[argmax_overlaps, 4]
labels_cat = labels_cat * (1.0 - F.less_equal(F.abs(labels), 1e-5))
ignore_mask = F.less_equal(F.abs(labels + 1), 1e-5)
labels_cat = labels_cat * (1 - ignore_mask) - ignore_mask
# assign low_quality boxes
if self.cfg.allow_low_quality:
gt_argmax_overlaps = F.argmax(overlaps, axis=0)
labels_cat = labels_cat.set_ai(gt_boxes[:, 4])[gt_argmax_overlaps]
matched_low_bbox_targets = self.box_coder.encode(
anchors.ai[gt_argmax_overlaps, :], gt_boxes[:, :4]
)
bbox_targets = bbox_targets.set_ai(matched_low_bbox_targets)[
gt_argmax_overlaps, :
]
labels_cat_list.append(F.add_axis(labels_cat, 0))
bbox_targets_list.append(F.add_axis(bbox_targets, 0))
return (
F.zero_grad(F.concat(labels_cat_list, axis=0)),
F.zero_grad(F.concat(bbox_targets_list, axis=0)),
)
def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45):
box_corner = F.zeros_like(prediction)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
# If none are remaining => process next image
if not image_pred.shape[0]:
continue
# Get score and class with highest confidence
class_conf = F.max(image_pred[:, 5 : 5 + num_classes], 1, keepdims=True)
class_pred = F.argmax(image_pred[:, 5 : 5 + num_classes], 1, keepdims=True)
class_conf_squeeze = F.squeeze(class_conf)
conf_mask = image_pred[:, 4] * class_conf_squeeze >= conf_thre
detections = F.concat((image_pred[:, :5], class_conf, class_pred), 1)
detections = detections[conf_mask]
if not detections.shape[0]:
continue
nms_out_index = F.vision.nms(
detections[:, :4], detections[:, 4] * detections[:, 5], nms_thre,
)
detections = detections[nms_out_index]
if output[i] is None:
output[i] = detections
else:
output[i] = F.concat((output[i], detections))
return output
def train_generator_batch(optical, sar, label, *, opt, netG):
netG.train()
cls_score, offsets, ctr_score = netG(sar, optical)
loss, loss_cls, loss_reg, loss_ctr = netG.loss(cls_score, offsets,
ctr_score, label)
opt.backward(loss)
if dist.is_distributed():
# do all reduce mean
pass
# performance in the training data
B, _, _, _ = cls_score.shape
cls_score = F.sigmoid(cls_score) # * ctr_score
cls_score = cls_score.reshape(B, -1)
# find the max
max_id = F.argmax(cls_score, axis=1) # (B, )
pred_box = get_box(netG.fm_ctr, offsets) # (B,4,H,W)
pred_box = pred_box.reshape(B, 4, -1)
output = []
for i in range(B):
output.append(F.add_axis(pred_box[i, :, max_id[i]], axis=0)) # (1, 4)
output = F.concat(output, axis=0) # (B, 4)
return [
loss_cls, loss_reg, loss_ctr,
F.norm(output[:, 0:2] - label[:, 0:2], p=2, axis=1).mean()
]
def train_generator_batch(optical, sar, label, *, opt, netG):
netG.train()
cls_score, offsets, ctr_score = netG(sar,
optical) #通过head得到分类分数、偏移和centerscore
loss, loss_cls, loss_reg, loss_ctr = netG.loss(cls_score, offsets,
ctr_score,
label) #调用Loss函数计算loss
opt.backward(loss)
if dist.is_distributed():
# do all reduce mean
pass
# performance in the training data 在训练数据上的性能
B, _, _, _ = cls_score.shape #B,1,37,37
cls_score = cls_score.reshape(B, -1)
# find the max
max_id = F.argmax(cls_score, axis=1) # (B, ) 找到最大的index
pred_box = get_box(netG.fm_ctr, offsets) # (B,4,H,W) 得到左上角点和右下角点
pred_box = pred_box.reshape(B, 4, -1) #(B,4,H*W)
output = []
for i in range(B): #找到每个图片中 预测出来的最大位置的bbox框 存入output
output.append(F.add_axis(pred_box[i, :, max_id[i]], axis=0)) # (1, 4)
output = F.concat(output, axis=0) # (B, 4)
return [
loss_cls, loss_reg,
F.norm(output[:, 0:2] - label[:, 0:2], p=2, axis=1).mean()
] #只计算左上角点的差距
def fpn_anchor_target_opr_core_impl(gt_boxes,
im_info,
anchors,
allow_low_quality_matches=True):
ignore_label = config.ignore_label
# get the gt boxes
valid_gt_boxes = gt_boxes[:im_info[5], :]
non_ignore_mask = valid_gt_boxes[:, -1] > 0
non_ignore_inds = mask_to_inds(non_ignore_mask)
valid_gt_boxes = valid_gt_boxes.ai[non_ignore_inds]
# compute the iou matrix
overlaps = box_overlap_opr(anchors, valid_gt_boxes[:, :4])
# match the dtboxes
a_shp0 = anchors.shape[0]
max_overlaps = F.max(overlaps, axis=1)
argmax_overlaps = F.argmax(overlaps, axis=1)
# all ignore
labels = mge.ones(a_shp0).astype(np.int32) * ignore_label
# set negative ones
labels = labels * (max_overlaps >= config.rpn_negative_overlap)
# set positive ones
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
const_one = mge.tensor(1.0)
if allow_low_quality_matches:
# match the max gt
gt_max_overlaps = F.max(overlaps, axis=0)
gt_argmax_overlaps = F.argmax(overlaps, axis=0)
g_shp0 = valid_gt_boxes.shapeof()[0]
gt_id = F.linspace(0, g_shp0 - 1, g_shp0).astype(np.int32)
argmax_overlaps = argmax_overlaps.set_ai(gt_id)[gt_argmax_overlaps]
max_overlaps = max_overlaps.set_ai(
const_one.broadcast(g_shp0))[gt_argmax_overlaps]
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
# set positive ones
fg_mask_ind = mask_to_inds(fg_mask)
labels = labels.set_ai(const_one.broadcast(
fg_mask_ind.shapeof()))[fg_mask_ind]
# compute the targets
bbox_targets = bbox_transform_opr(anchors,
valid_gt_boxes.ai[argmax_overlaps, :4])
if config.rpn_bbox_normalize_targets:
std_opr = mge.tensor(config.bbox_normalize_stds[None, :])
mean_opr = mge.tensor(config.bbox_normalize_means[None, :])
minus_opr = mean_opr / std_opr
bbox_targets = bbox_targets / std_opr - minus_opr
return labels, bbox_targets
def get_ground_truth(self, rpn_rois, im_info, gt_boxes):
if not self.training:
return rpn_rois, None, None
return_rois = []
return_labels = []
return_bbox_targets = []
# get per image proposals and gt_boxes
for bid in range(gt_boxes.shape[0]):
num_valid_boxes = im_info[bid, 4].astype("int32")
gt_boxes_per_img = gt_boxes[bid, :num_valid_boxes, :]
batch_inds = F.full((gt_boxes_per_img.shape[0], 1), bid)
gt_rois = F.concat([batch_inds, gt_boxes_per_img[:, :4]], axis=1)
batch_roi_mask = rpn_rois[:, 0] == bid
# all_rois : [batch_id, x1, y1, x2, y2]
all_rois = F.concat([rpn_rois[batch_roi_mask], gt_rois])
overlaps = layers.get_iou(all_rois[:, 1:5], gt_boxes_per_img)
max_overlaps = overlaps.max(axis=1)
gt_assignment = F.argmax(overlaps, axis=1).astype("int32")
labels = gt_boxes_per_img[gt_assignment, 4]
# ---------------- get the fg/bg labels for each roi ---------------#
fg_mask = (max_overlaps >= self.cfg.fg_threshold) & (labels >= 0)
bg_mask = ((max_overlaps >= self.cfg.bg_threshold_low)
& (max_overlaps < self.cfg.bg_threshold_high))
num_fg_rois = int(self.cfg.num_rois * self.cfg.fg_ratio)
fg_inds_mask = layers.sample_mask_from_labels(
fg_mask, num_fg_rois, 1)
num_bg_rois = int(self.cfg.num_rois - fg_inds_mask.sum())
bg_inds_mask = layers.sample_mask_from_labels(
bg_mask, num_bg_rois, 1)
labels = labels * fg_inds_mask
keep_mask = fg_inds_mask + bg_inds_mask
_, keep_inds = F.cond_take(keep_mask == 1, keep_mask)
# Add next line to avoid memory exceed
keep_inds = keep_inds[:min(self.cfg.num_rois, keep_inds.shape[0])]
labels = labels[keep_inds].astype("int32")
rois = all_rois[keep_inds]
target_boxes = gt_boxes_per_img[gt_assignment[keep_inds], :4]
bbox_targets = self.box_coder.encode(rois[:, 1:5], target_boxes)
bbox_targets = bbox_targets.reshape(-1, 4)
return_rois.append(rois)
return_labels.append(labels)
return_bbox_targets.append(bbox_targets)
return (
F.concat(return_rois, axis=0).detach(),
F.concat(return_labels, axis=0).detach(),
F.concat(return_bbox_targets, axis=0).detach(),
)
def per_level_gt(self,
gt_boxes,
im_info,
anchors,
allow_low_quality_matches=True):
ignore_label = self.cfg.ignore_label
# get the gt boxes
valid_gt_boxes = gt_boxes[:im_info[4], :]
# compute the iou matrix
overlaps = layers.get_iou(anchors, valid_gt_boxes[:, :4])
# match the dtboxes
a_shp0 = anchors.shape[0]
max_overlaps = F.max(overlaps, axis=1)
argmax_overlaps = F.argmax(overlaps, axis=1)
# all ignore
labels = mge.ones(a_shp0).astype("int32") * ignore_label
# set negative ones
labels = labels * (max_overlaps >= self.cfg.rpn_negative_overlap)
# set positive ones
fg_mask = max_overlaps >= self.cfg.rpn_positive_overlap
const_one = mge.tensor(1.0)
if allow_low_quality_matches:
# make sure that max iou of gt matched
gt_argmax_overlaps = F.argmax(overlaps, axis=0)
num_valid_boxes = valid_gt_boxes.shapeof(0)
gt_id = F.linspace(0, num_valid_boxes - 1,
num_valid_boxes).astype("int32")
argmax_overlaps = argmax_overlaps.set_ai(gt_id)[gt_argmax_overlaps]
max_overlaps = max_overlaps.set_ai(
const_one.broadcast(num_valid_boxes))[gt_argmax_overlaps]
fg_mask = max_overlaps >= self.cfg.rpn_positive_overlap
# set positive ones
_, fg_mask_ind = F.cond_take(fg_mask == 1, fg_mask)
labels = labels.set_ai(const_one.broadcast(
fg_mask_ind.shapeof(0)))[fg_mask_ind]
# compute the targets
bbox_targets = self.box_coder.encode(
anchors, valid_gt_boxes.ai[argmax_overlaps, :4])
return labels, bbox_targets
def test_generator_batch(optical, sar, *, netG):
netG.eval()
cls_score, offsets, ctr_score = netG(
sar, optical) # [B,1,19,19] [B,2,19,19] [B,1,19,19]
B, _, _, _ = cls_score.shape
# 加权
# cls_score = cls_score * ctr_score
cls_score = cls_score.reshape(B, -1) # [B,19*19]
# find the max
max_id = F.argmax(cls_score, axis=1) # (B, )
pred_box = get_box(netG.fm_ctr, offsets) # (B,4,H,W)
pred_box = pred_box.reshape(B, 4, -1)
output = []
for i in range(B):
output.append(F.add_axis(pred_box[i, :, max_id[i]], axis=0)) # (1, 4)
return F.concat(output, axis=0) # [B,4]
def __call__(self, matrix):
"""
matrix(tensor): A two dim tensor with shape of (N, M). N is number of GT-boxes,
while M is the number of anchors in detection.
"""
assert len(matrix.shape) == 2
max_scores = matrix.max(axis=0)
match_indices = F.argmax(matrix, axis=0)
# default ignore label: -1
labels = F.full_like(match_indices, -1)
for label, low, high in zip(self.labels, self.thresholds[:-1],
self.thresholds[1:]):
mask = (max_scores >= low) & (max_scores < high)
labels[mask] = label
if self.allow_low_quality_matches:
mask = (matrix == F.max(matrix, axis=1,
keepdims=True)).sum(axis=0) > 0
labels[mask] = 1
return match_indices, labels
def accuracy(out, labels):
outputs = F.argmax(out, axis=1)
return F.sum(outputs == labels)
def fwd(data):
return F.argmax(data), F.argmin(data)
def get_ground_truth(self, anchors_list, batched_gt_boxes,
batched_num_gts):
labels_list = []
offsets_list = []
ctrness_list = []
all_level_anchors = F.concat(anchors_list, axis=0)
for bid in range(batched_gt_boxes.shape[0]):
gt_boxes = batched_gt_boxes[bid, :batched_num_gts[bid]]
ious = []
candidate_idxs = []
base = 0
for stride, anchors_i in zip(self.cfg.stride, anchors_list):
ious.append(
layers.get_iou(
gt_boxes[:, :4],
F.concat([
anchors_i - stride * self.cfg.anchor_scale / 2,
anchors_i + stride * self.cfg.anchor_scale / 2,
],
axis=1)))
gt_centers = (gt_boxes[:, :2] + gt_boxes[:, 2:4]) / 2
distances = F.sqrt(
F.sum((F.expand_dims(gt_centers, axis=1) - anchors_i)**2,
axis=2))
_, topk_idxs = F.topk(distances, self.cfg.anchor_topk)
candidate_idxs.append(base + topk_idxs)
base += anchors_i.shape[0]
ious = F.concat(ious, axis=1)
candidate_idxs = F.concat(candidate_idxs, axis=1)
candidate_ious = F.gather(ious, 1, candidate_idxs)
ious_thr = (F.mean(candidate_ious, axis=1, keepdims=True) +
F.std(candidate_ious, axis=1, keepdims=True))
is_foreground = F.scatter(
F.zeros(ious.shape), 1, candidate_idxs,
F.ones(candidate_idxs.shape)).astype(bool) & (ious >= ious_thr)
is_in_boxes = F.min(self.point_coder.encode(
all_level_anchors, F.expand_dims(gt_boxes[:, :4], axis=1)),
axis=2) > 0
ious[~is_foreground] = -1
ious[~is_in_boxes] = -1
match_indices = F.argmax(ious, axis=0)
gt_boxes_matched = gt_boxes[match_indices]
anchor_max_iou = F.indexing_one_hot(ious, match_indices, axis=0)
labels = gt_boxes_matched[:, 4].astype(np.int32)
labels[anchor_max_iou == -1] = 0
offsets = self.point_coder.encode(all_level_anchors,
gt_boxes_matched[:, :4])
left_right = offsets[:, [0, 2]]
top_bottom = offsets[:, [1, 3]]
ctrness = F.sqrt(
F.clip(F.min(left_right, axis=1) / F.max(left_right, axis=1),
lower=0) *
F.clip(F.min(top_bottom, axis=1) / F.max(top_bottom, axis=1),
lower=0))
labels_list.append(labels)
offsets_list.append(offsets)
ctrness_list.append(ctrness)
return (
F.stack(labels_list, axis=0).detach(),
F.stack(offsets_list, axis=0).detach(),
F.stack(ctrness_list, axis=0).detach(),
)
def get_ground_truth(self, rpn_rois, im_info, gt_boxes):
if not self.training:
return rpn_rois, None, None
return_rois = []
return_labels = []
return_bbox_targets = []
# get per image proposals and gt_boxes
for bid in range(self.cfg.batch_per_gpu):
num_valid_boxes = im_info[bid, 4]
gt_boxes_per_img = gt_boxes[bid, :num_valid_boxes, :]
batch_inds = mge.ones((gt_boxes_per_img.shapeof(0), 1)) * bid
# if config.proposal_append_gt:
gt_rois = F.concat([batch_inds, gt_boxes_per_img[:, :4]], axis=1)
batch_roi_mask = rpn_rois[:, 0] == bid
_, batch_roi_inds = F.cond_take(batch_roi_mask == 1, batch_roi_mask)
# all_rois : [batch_id, x1, y1, x2, y2]
all_rois = F.concat([rpn_rois.ai[batch_roi_inds], gt_rois])
overlaps_normal, overlaps_ignore = layers.get_iou(
all_rois[:, 1:5], gt_boxes_per_img, return_ignore=True,
)
max_overlaps_normal = overlaps_normal.max(axis=1)
gt_assignment_normal = F.argmax(overlaps_normal, axis=1)
max_overlaps_ignore = overlaps_ignore.max(axis=1)
gt_assignment_ignore = F.argmax(overlaps_ignore, axis=1)
ignore_assign_mask = (max_overlaps_normal < self.cfg.fg_threshold) * (
max_overlaps_ignore > max_overlaps_normal
)
max_overlaps = (
max_overlaps_normal * (1 - ignore_assign_mask)
+ max_overlaps_ignore * ignore_assign_mask
)
gt_assignment = (
gt_assignment_normal * (1 - ignore_assign_mask)
+ gt_assignment_ignore * ignore_assign_mask
)
gt_assignment = gt_assignment.astype("int32")
labels = gt_boxes_per_img.ai[gt_assignment, 4]
# ---------------- get the fg/bg labels for each roi ---------------#
fg_mask = (max_overlaps >= self.cfg.fg_threshold) * (
labels != self.cfg.ignore_label
)
bg_mask = (max_overlaps < self.cfg.bg_threshold_high) * (
max_overlaps >= self.cfg.bg_threshold_low
)
num_fg_rois = self.cfg.num_rois * self.cfg.fg_ratio
fg_inds_mask = self._bernoulli_sample_masks(fg_mask, num_fg_rois, 1)
num_bg_rois = self.cfg.num_rois - fg_inds_mask.sum()
bg_inds_mask = self._bernoulli_sample_masks(bg_mask, num_bg_rois, 1)
labels = labels * fg_inds_mask
keep_mask = fg_inds_mask + bg_inds_mask
_, keep_inds = F.cond_take(keep_mask == 1, keep_mask)
# Add next line to avoid memory exceed
keep_inds = keep_inds[: F.minimum(self.cfg.num_rois, keep_inds.shapeof(0))]
# labels
labels = labels.ai[keep_inds].astype("int32")
rois = all_rois.ai[keep_inds]
target_boxes = gt_boxes_per_img.ai[gt_assignment.ai[keep_inds], :4]
bbox_targets = self.box_coder.encode(rois[:, 1:5], target_boxes)
bbox_targets = bbox_targets.reshape(-1, 4)
return_rois.append(rois)
return_labels.append(labels)
return_bbox_targets.append(bbox_targets)
return (
F.zero_grad(F.concat(return_rois, axis=0)),
F.zero_grad(F.concat(return_labels, axis=0)),
F.zero_grad(F.concat(return_bbox_targets, axis=0)),
)
def run_argmax():
x = F.zeros((100, 100))
x[:] = -float("inf")
idxs = F.argmax(x, axis=0)
return idxs
import cv2
import os
import numpy as np
star = time.time()
le_net = Unet(1, 4)
state_dict = mge.load("unet.mge")
le_net.load_state_dict(state_dict)
le_net.eval() # 设置为测试模式
image = mge.tensor(dtype="float32")
path = "./data/train_l/image"
file_name = os.listdir(path)
for step, (name) in enumerate(file_name):
image_ = cv2.imread(os.path.join(path, name), cv2.IMREAD_GRAYSCALE)
image_ = image_ / 255.
image_ = image_.reshape(1, 1, 256, 256).astype(np.float32)
image.set_value(image_)
logits = le_net(image)
predicted = F.argmax(logits, axis=1)
predicted = predicted.numpy()
predicted = predicted.transpose(1, 2, 0)
cv2.imwrite("./result/%s" % name, predicted * 80)
print("end using time:%.4f" % (time.time() - star))
