def boxes_to_masks(boxes, h, w, padding=0.0):
n = boxes.shape[0]
boxes = boxes
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
b_w = x2 - x1
b_h = y2 - y1
x1 = jt.clamp(x1 - 1 - b_w * padding, min_v=0)
x2 = jt.clamp(x2 + 1 + b_w * padding, max_v=w)
y1 = jt.clamp(y1 - 1 - b_h * padding, min_v=0)
y2 = jt.clamp(y2 + 1 + b_h * padding, max_v=h)
rows = jt.arange(w, dtype=x1.dtype).view(1, 1, -1).expand((n, h, w))
cols = jt.arange(h, dtype=x1.dtype).view(1, -1, 1).expand((n, h, w))
masks_left = rows >= x1.view(-1, 1, 1)
masks_right = rows < x2.view(-1, 1, 1)
masks_up = cols >= y1.view(-1, 1, 1)
masks_down = cols < y2.view(-1, 1, 1)
masks = masks_left * masks_right * masks_up * masks_down
return masks
def crop(masks, boxes, padding: int = 1):
"""
"Crop" predicted masks by zeroing out everything not in the predicted bbox.
Vectorized by Chong (thanks Chong).
Args:
- masks should be a size [h, w, n] tensor of masks
- boxes should be a size [n, 4] tensor of bbox coords in relative point form
"""
h, w, n = masks.shape
x1, x2 = sanitize_coordinates(boxes[:, 0],
boxes[:, 2],
w,
padding,
cast=False)
y1, y2 = sanitize_coordinates(boxes[:, 1],
boxes[:, 3],
h,
padding,
cast=False)
rows = jt.arange(w, dtype=x1.dtype).view(1, -1, 1).expand((h, w, n))
cols = jt.arange(h, dtype=x1.dtype).view(-1, 1, 1).expand((h, w, n))
masks_left = rows >= x1.view(1, 1, -1)
masks_right = rows < x2.view(1, 1, -1)
masks_up = cols >= y1.view(1, 1, -1)
masks_down = cols < y2.view(1, 1, -1)
crop_mask = masks_left * masks_right * masks_up * masks_down
return masks * crop_mask.float()
def compute_locations_per_level(self, h, w, stride):
shifts_x = jt.arange(0, w * stride, step=stride, dtype='float32')
shifts_y = jt.arange(0, h * stride, step=stride, dtype='float32')
shift_y, shift_x = jt.meshgrid(shifts_y, shifts_x)
shift_x = shift_x.reshape(-1)
shift_y = shift_y.reshape(-1)
locations = jt.stack((shift_x, shift_y), dim=1) + stride // 2
return locations
def execute(self, x, boxes):
"""
Arguments:
x (Tensor): the mask logits
boxes (list[BoxList]): bounding boxes that are used as
reference, one for each image
Returns:
results (list[BoxList]): one BoxList for each image, containing
the extra field mask
"""
mask_prob = x.sigmoid()
# select masks coresponding to the predicted classes
num_masks = x.shape[0]
labels = [bbox.get_field("labels") for bbox in boxes]
labels = jt.contrib.concat(labels, dim=0)
index = jt.arange(num_masks)
mask_prob = mask_prob[index, labels].unsqueeze(1)
boxes_per_image = [len(box) for box in boxes]
mask_prob = mask_prob.split(boxes_per_image, dim=0)
if self.masker:
mask_prob = self.masker(mask_prob, boxes)
results = []
for prob, box in zip(mask_prob, boxes):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
bbox.add_field("mask", prob)
results.append(bbox)
return results
def __call__(self, proposals, mask_logits, targets):
"""
Arguments:
proposals (list[BoxList])
mask_logits (Tensor)
targets (list[BoxList])
Return:
mask_loss (Tensor): scalar tensor containing the loss
"""
labels, mask_targets, mask_ratios = self.prepare_targets(
proposals, targets)
labels = cat(labels, dim=0)
mask_targets = cat(mask_targets, dim=0)
positive_inds = jt.nonzero(labels > 0).squeeze(1)
labels_pos = labels[positive_inds]
# accept empty tensors, so handle it separately
if mask_targets.numel() == 0:
if not self.maskiou_on:
return mask_logits.sum() * 0
else:
selected_index = jt.arange(mask_logits.shape[0])
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h,
mask_w)
return mask_logits.sum() * 0, selected_mask, labels, None
if self.maskiou_on:
mask_ratios = cat(mask_ratios, dim=0)
value_eps = 1e-10 * jt.ones((mask_targets.shape[0], ))
mask_ratios = jt.maximum(mask_ratios, value_eps)
pred_masks = mask_logits[positive_inds, labels_pos]
pred_masks[:] = pred_masks > 0
mask_targets_full_area = mask_targets.sum(
dims=[1, 2]) / mask_ratios
mask_ovr = pred_masks * mask_targets
mask_ovr_area = mask_ovr.sum(dims=[1, 2])
mask_union_area = pred_masks.sum(
dims=[1, 2]) + mask_targets_full_area - mask_ovr_area
value_1 = jt.ones((pred_masks.shape[0], ))
value_0 = jt.zeros((pred_masks.shape[0], ))
mask_union_area = jt.maximum(mask_union_area, value_1)
mask_ovr_area = jt.maximum(mask_ovr_area, value_0)
maskiou_targets = mask_ovr_area / mask_union_area
binary_cross_entropy_with_logits = nn.BCEWithLogitsLoss()
mask_loss = binary_cross_entropy_with_logits(
mask_logits[positive_inds, labels_pos], mask_targets)
if not self.maskiou_on:
return mask_loss
else:
selected_index = jt.index((mask_logits.shape[0], ), dim=0)
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h, mask_w)
selected_mask = selected_mask.sigmoid()
return mask_loss, selected_mask, labels, maskiou_targets
def test_normal(self):
from jittor import init
n = 10000
r = 0.155
a = init.gauss([n], "float32", 1, 3)
data = a.data
assert (np.abs((data < (1 - 3)).mean() - r) < 0.1)
assert (np.abs((data < (1)).mean() - 0.5) < 0.1)
assert (np.abs((data < (1 + 3)).mean() - (1 - r)) < 0.1)
np_res = np.random.normal(1, 0.1, (100, 100))
jt_res = jt.normal(1., 0.1, (100, 100))
assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1)
assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1)
np_res = torch.normal(torch.arange(1., 10000.), 1)
jt_res = jt.normal(jt.arange(1, 10000), 1)
assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1)
assert (np.abs(np_res.std() - jt_res.data.std()) < 1)
np_res = np.random.randn(100, 100)
jt_res = jt.randn(100, 100)
assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1)
assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1)
np_res = np.random.rand(100, 100)
jt_res = jt.rand(100, 100)
assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1)
assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1)
def add_self_loops(edge_index,
edge_weight: Optional[Var] = None,
fill_value: float = 1.,
num_nodes: Optional[int] = None):
r"""Adds a self-loop :math:`(i,i) \in \mathcal{E}` to every node
:math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`.
In case the graph is weighted, self-loops will be added with edge weights
denoted by :obj:`fill_value`.
Args:
edge_index (Var int32): The edge indices.
edge_weight (Var, optional): One-dimensional edge weights.
(default: :obj:`None`)
fill_value (float, optional): If :obj:`edge_weight` is not :obj:`None`,
will add self-loops with edge weights of :obj:`fill_value` to the
graph. (default: :obj:`1.`)
num_nodes (int, optional): The number of nodes, *i.e.*
:obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
:rtype: (:class:`Var int32`, :class:`Var`)
"""
N = maybe_num_nodes(edge_index, num_nodes)
loop_index = jt.arange(0, N, dtype=Var.int32)
loop_index = loop_index.unsqueeze(0).repeat(2, 1)
if edge_weight is not None:
assert edge_weight.numel() == edge_index.size(1)
loop_weight = init.constant((N, ), edge_weight.dtype, fill_value)
edge_weight = jt.concat([edge_weight, loop_weight], dim=0)
edge_index = jt.concat([edge_index, loop_index], dim=1)
return edge_index, edge_weight
def grid_anchors(self, grid_sizes):
anchors = []
for size, stride, base_anchors in zip(grid_sizes, self.strides,
self.cell_anchors):
grid_height, grid_width = size
shifts_x = jt.arange(0, grid_width * stride, step=stride).float32()
shifts_y = jt.arange(0, grid_height * stride,
step=stride).float32()
shift_y, shift_x = jt.meshgrid(shifts_y, shifts_x)
shift_x = shift_x.reshape(-1)
shift_y = shift_y.reshape(-1)
shifts = jt.stack((shift_x, shift_y, shift_x, shift_y), dim=1)
anchors.append((shifts.reshape(-1, 1, 4) +
base_anchors.reshape(1, -1, 4)).reshape(-1, 4))
return anchors
def execute(self, x):
assert x.ndim == 3
# interp = torch.lerp(self.avg_latent, x, self.threshold)
interp = self.avg_latent + self.threshold * (x - self.avg_latent)
# do_trunc = (torch.arange(x.size(1)) < self.max_layer).view(1, -1, 1).to(x.device)
do_trunc = (jt.arange(x.size(1)) < self.max_layer).view(1, -1, 1)
# return torch.where(do_trunc, interp, x)
return do_trunc * interp + (1 - do_trunc) * x
def fast_nms(self,
boxes,
masks,
scores,
iou_threshold: float = 0.5,
top_k: int = 200,
second_threshold: bool = False):
idx, scores = scores.argsort(1, descending=True)
idx = idx[:, :top_k]
scores = scores[:, :top_k]
num_classes, num_dets = idx.shape
boxes = boxes[idx.view(-1)].view(num_classes, num_dets, 4)
masks = masks[idx.view(-1)].view(num_classes, num_dets, -1)
iou = jaccard(boxes, boxes)
iou = iou.triu_(diagonal=1)
iou_max = iou.max(dim=1)
# Now just filter out the ones higher than the threshold
keep = (iou_max <= iou_threshold)
# We should also only keep detections over the confidence threshold, but at the cost of
# maxing out your detection count for every image, you can just not do that. Because we
# have such a minimal amount of computation per detection (matrix mulitplication only),
# this increase doesn't affect us much (+0.2 mAP for 34 -> 33 fps), so we leave it out.
# However, when you implement this in your method, you should do this second threshold.
if second_threshold:
keep *= (scores > self.conf_thresh)
# Assign each kept detection to its corresponding class
classes = jt.arange(num_classes).unsqueeze(1).expand_as(keep)
classes = classes[keep]
boxes = boxes[keep]
masks = masks[keep]
scores = scores[keep]
#print('keep finish')
# Only keep the top cfg.max_num_detections highest scores across all classes
idx, scores = jt.argsort(scores, dim=0, descending=True)
#print('argsort finish')
idx = idx[:cfg.max_num_detections]
scores = scores[:cfg.max_num_detections]
classes = classes[idx]
boxes = boxes[idx]
masks = masks[idx]
'''
scores = scores[:cfg.max_num_detections]
classes = classes[:cfg.max_num_detections]
boxes = boxes[:cfg.max_num_detections]
masks = masks[:cfg.max_num_detections]
'''
return boxes, masks, classes, scores
def remove_isolated_nodes(edge_index, edge_attr=None, num_nodes=None):
r"""Removes the isolated nodes from the graph given by :attr:`edge_index`
with optional edge attributes :attr:`edge_attr`.
In addition, returns a mask of shape :obj:`[num_nodes]` to manually filter
out isolated node features later on.
Self-loops are preserved for non-isolated nodes.
Args:
edge_index (Var int32): The edge indices.
edge_attr (Var, optional): Edge weights or multi-dimensional
edge features. (default: :obj:`None`)
num_nodes (int, optional): The number of nodes, *i.e.*
:obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
:rtype: (Var int32, Var, Var bool)
"""
num_nodes = maybe_num_nodes(edge_index, num_nodes)
out = segregate_self_loops(edge_index, edge_attr)
edge_index, edge_attr, loop_edge_index, loop_edge_attr = out
mask = jt.zeros((num_nodes), dtype=Var.bool)
mask[edge_index.view(-1)] = 1
assoc = jt.full((num_nodes, ), -1, dtype=Var.int32)
assoc[mask] = jt.arange(mask.sum())
edge_index = assoc[edge_index]
loop_mask = jt.zeros_like(mask)
loop_mask[loop_edge_index[0]] = 1
loop_mask = loop_mask & mask
loop_assoc = jt.full_like(assoc, -1)
loop_assoc[loop_edge_index[0]] = jt.arange(loop_edge_index.size(1))
loop_idx = loop_assoc[loop_mask]
loop_edge_index = assoc[loop_edge_index[:, loop_idx]]
edge_index = jt.concat([edge_index, loop_edge_index], dim=1)
if edge_attr is not None:
loop_edge_attr = loop_edge_attr[loop_idx]
edge_attr = jt.concat([edge_attr, loop_edge_attr], dim=0)
return edge_index, edge_attr, mask
def crop_by_box(masks, box, padding=0.0):
n, h, w = masks.size()
b_w = box[:, 2] - box[:, 0]
b_h = box[:, 3] - box[:, 1]
x1 = jt.clamp(box[:, 0] - b_w * padding - 1, min_v=0)
x2 = jt.clamp(box[:, 2] + b_w * padding + 1, max_v=w - 1)
y1 = jt.clamp(box[:, 1] - b_h * padding - 1, min_v=0)
y2 = jt.clamp(box[:, 3] + b_h * padding + 1, max_v=h - 1)
rows = jt.arange(w, dtype=x1.dtype).view(1, 1, -1).expand((n, h, w))
cols = jt.arange(h, dtype=x1.dtype).view(1, -1, 1).expand((n, h, w))
masks_left = rows >= x1.view(n, 1, 1)
masks_right = rows < x2.view(n, 1, 1)
masks_up = cols >= y1.view(n, 1, 1)
masks_down = cols < y2.view(n, 1, 1)
crop_mask = masks_left * masks_right * masks_up * masks_down
return masks * crop_mask.float(), crop_mask
def sigmoid_focal_loss(logits, targets, gamma, alpha):
num_classes = logits.shape[1]
dtype = targets.dtype
class_range = jt.arange(1, num_classes + 1, dtype=dtype).unsqueeze(0)
t = targets.unsqueeze(1)
p = logits.sigmoid()
term1 = (1 - p)**gamma * jt.log(p)
term2 = p**gamma * jt.log(1 - p)
return -(t == class_range).float() * term1 * alpha - (
(t != class_range) * (t >= 0)).float() * term2 * (1 - alpha)
def traditional_nms(self,
boxes,
masks,
scores,
iou_threshold=0.5,
conf_thresh=0.05):
import pyximport
pyximport.install(setup_args={"include_dirs": np.get_include()},
reload_support=True)
from utils.cython_nms import nms as cnms
num_classes = scores.shape[0]
idx_lst = []
cls_lst = []
scr_lst = []
# Multiplying by max_size is necessary because of how cnms computes its area and intersections
boxes = boxes * cfg.max_size
for _cls in range(num_classes):
cls_scores = scores[_cls]
conf_mask = cls_scores > conf_thresh
idx = jt.arange(cls_scores.shape[0], device=boxes.device)
cls_scores = cls_scores[conf_mask]
idx = idx[conf_mask]
if cls_scores.shape[0] == 0:
continue
preds = jt.contrib.concat([boxes[conf_mask], cls_scores[:, None]],
dim=1).numpy()
keep = cnms(preds, iou_threshold)
keep = jt.array(keep).int64()
idx_lst.append(idx[keep])
cls_lst.append(keep * 0 + _cls)
scr_lst.append(cls_scores[keep])
idx = jt.contrib.concat(idx_lst, dim=0)
classes = jt.contrib.concat(cls_lst, dim=0)
scores = jt.contrib.concat(scr_lst, dim=0)
scores, idx2 = scores.sort(0, descending=True)
idx2 = idx2[:cfg.max_num_detections]
scores = scores[:cfg.max_num_detections]
idx = idx[idx2]
classes = classes[idx2]
# Undo the multiplication above
return boxes[idx] / cfg.max_size, masks[idx], classes, scores
def random_subsample(pcd, n_points=2048):
"""
Args:
pcd: (B, N, 3)
returns:
new_pcd: (B, n_points, 3)
"""
b, n, _ = pcd.shape
batch_idx = jittor.arange(b).reshape((-1, 1)).repeat(1, n_points)
idx = jittor.concat([jittor.randperm(n)[:n_points].reshape((1, -1)) for i in range(b)], 0)
return pcd[batch_idx, idx, :]
def grad(self, grad):
x, numangle, numrho = self.save_vars
cuda_src_backward = csb.replace('#numangle', str(numangle))
cuda_src_backward = cuda_src_backward.replace('#numrho', str(numrho))
irho = int((h * h + w * w)**0.5 + 1) / float((numrho - 1))
itheta = 3.14159265358979323846 / numangle
angle = jt.arange(numangle) * itheta
tabCos = angle.cos() / irho
tabSin = angle.sin() / irho
return jt.code([x.shape], [x.dtype], [x, grad, tabCos, tabSin],
cuda_src=cuda_src_backward)
def _forward_train(self,features,img_size,boxes,labels):
N = features.shape[0]
rpn_locs, rpn_scores, rois, roi_indices, anchor = self.rpn(features, img_size)
sample_rois = []
gt_roi_locs = []
gt_roi_labels = []
sample_roi_indexs = []
gt_rpn_locs = []
gt_rpn_labels = []
for i in range(N):
index = jt.where(roi_indices == i)[0]
roi = rois[index,:]
box = boxes[i]
label = labels[i]
sample_roi, gt_roi_loc, gt_roi_label = self.proposal_target_creator(roi,box,label)
sample_roi_index = i*jt.ones((sample_roi.shape[0],))
sample_rois.append(sample_roi)
gt_roi_labels.append(gt_roi_label)
gt_roi_locs.append(gt_roi_loc)
sample_roi_indexs.append(sample_roi_index)
gt_rpn_loc, gt_rpn_label = self.anchor_target_creator(box,anchor,img_size)
gt_rpn_locs.append(gt_rpn_loc)
gt_rpn_labels.append(gt_rpn_label)
sample_roi_indexs = jt.contrib.concat(sample_roi_indexs,dim=0)
sample_rois = jt.contrib.concat(sample_rois,dim=0)
roi_cls_loc, roi_score = self.head(features,sample_rois,sample_roi_indexs)
# ------------------ RPN losses -------------------#
rpn_locs = rpn_locs.reshape(-1,4)
rpn_scores = rpn_scores.reshape(-1,2)
gt_rpn_labels = jt.contrib.concat(gt_rpn_labels,dim=0)
gt_rpn_locs = jt.contrib.concat(gt_rpn_locs,dim=0)
rpn_loc_loss = _fast_rcnn_loc_loss(rpn_locs,gt_rpn_locs,gt_rpn_labels,self.rpn_sigma)
rpn_cls_loss = nn.cross_entropy_loss(rpn_scores[gt_rpn_labels>=0,:],gt_rpn_labels[gt_rpn_labels>=0])
# ------------------ ROI losses (fast rcnn loss) -------------------#
gt_roi_locs = jt.contrib.concat(gt_roi_locs,dim=0)
gt_roi_labels = jt.contrib.concat(gt_roi_labels,dim=0)
n_sample = roi_cls_loc.shape[0]
roi_cls_loc = roi_cls_loc.view(n_sample, np.prod(roi_cls_loc.shape[1:]).item()//4, 4)
roi_loc = roi_cls_loc[jt.arange(0, n_sample).int32(), gt_roi_labels]
roi_loc_loss = _fast_rcnn_loc_loss(roi_loc,gt_roi_locs,gt_roi_labels,self.roi_sigma)
roi_cls_loss = nn.cross_entropy_loss(roi_score, gt_roi_labels)
losses = [rpn_loc_loss, rpn_cls_loss, roi_loc_loss, roi_cls_loss]
losses = losses + [sum(losses)]
return losses
def forward_for_single_feature_map(self, anchors, objectness, box_regression):
"""
Arguments:
anchors: list[BoxList]
objectness: tensor of size N, A, H, W
box_regression: tensor of size N, A * 4, H, W
"""
N, A, H, W = objectness.shape
# put in the same format as anchors
objectness = permute_and_flatten(objectness, N, A, 1, H, W).reshape(N, -1)
objectness = objectness.sigmoid()
box_regression = permute_and_flatten(box_regression, N, A, 4, H, W)
num_anchors = A * H * W
pre_nms_top_n = min(self.pre_nms_top_n, num_anchors)
objectness, topk_idx = objectness.topk(pre_nms_top_n, dim=1, sorted=True)
batch_idx = jt.arange(N).unsqueeze(1)
box_regression = box_regression[batch_idx,topk_idx]
image_shapes = [box.size for box in anchors]
concat_anchors = jt.contrib.concat([a.bbox for a in anchors], dim=0)
concat_anchors = concat_anchors.reshape(N, -1, 4)[batch_idx, topk_idx]
proposals = self.box_coder.decode(
box_regression.reshape(-1, 4), concat_anchors.reshape(-1, 4)
)
proposals = proposals.reshape(N, -1, 4)
result = []
for i in range(len(image_shapes)):
proposal, score, im_shape = proposals[i], objectness[i], image_shapes[i]
boxlist = BoxList(proposal, im_shape, mode="xyxy")
boxlist.add_field("objectness", score)
boxlist = boxlist.clip_to_image(remove_empty=False)
boxlist = remove_small_boxes(boxlist, self.min_size)
boxlist = boxlist_nms(
boxlist,
self.nms_thresh,
max_proposals=self.post_nms_top_n,
score_field="objectness",
)
result.append(boxlist)
return result
def index2d(src, idx):
"""
Indexes a tensor by a 2d index.
In effect, this does
out[i, j] = src[i, idx[i, j]]
Both src and idx should have the same size.
"""
offs = jt.arange((idx.shape[0])).unsqueeze(1).expand_as(idx)
idx = idx + offs * idx.shape[1]
return src.view(-1)[idx.view(-1)].view(idx.shape)
def execute(self, x, numangle, numrho):
n, c, h, w = x.shape
cuda_src_forward = csf.replace('#numangle', str(numangle))
cuda_src_forward = cuda_src_forward.replace('#numrho', str(numrho))
irho = int((h * h + w * w)**0.5 + 1) / float((numrho - 1))
itheta = 3.14159265358979323846 / numangle
angle = jt.arange(numangle) * itheta
tabCos = angle.cos() / irho
tabSin = angle.sin() / irho
output = jt.code([n, c, numangle, numrho],
x.dtype, [x, tabCos, tabSin],
cuda_src=cuda_src_forward)
self.save_vars = x, numangle, numrho
return output
def add_remaining_self_loops(edge_index,
edge_weight: Optional[Var] = None,
fill_value: float = 1.,
num_nodes: Optional[int] = None):
r"""Adds remaining self-loop :math:`(i,i) \in \mathcal{E}` to every node
:math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`.
In case the graph is weighted and already contains a few self-loops, only
non-existent self-loops will be added with edge weights denoted by
:obj:`fill_value`.
Args:
edge_index (Var int32): The edge indices.
edge_weight (Var, optional): One-dimensional edge weights.
(default: :obj:`None`)
fill_value (float, optional): If :obj:`edge_weight` is not :obj:`None`,
will add self-loops with edge weights of :obj:`fill_value` to the
graph. (default: :obj:`1.`)
num_nodes (int, optional): The number of nodes, *i.e.*
:obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
:rtype: (:class:`Var int32`, :class:`Var`)
"""
N = maybe_num_nodes(edge_index, num_nodes)
row, col = edge_index[0], edge_index[1]
mask = row != col
loop_index = jt.arange(0, N, dtype=row.dtype)
loop_index = loop_index.unsqueeze(0).repeat(2, 1)
edge_index = jt.concat([edge_index[:, mask], loop_index], dim=1)
if edge_weight is not None:
inv_mask = jt.logical_not(mask)
loop_weight = init.constant((N, ), edge_weight.dtype, fill_value)
remaining_edge_weight = edge_weight[inv_mask]
if remaining_edge_weight.numel() > 0:
loop_weight[row[inv_mask]] = remaining_edge_weight
edge_weight = jt.concat([edge_weight[mask], loop_weight], dim=0)
return edge_index, edge_weight
def test_scatter(self):
src = jt.arange(1, 11).reshape((2, 5))
index = jt.array([[0, 1, 2, 0]])
x = jt.zeros((3, 5), dtype=src.dtype).scatter_(0, index, src)
assert (x.data == [[1, 0, 0, 4, 0], [0, 2, 0, 0, 0], [0, 0, 3, 0,
0]]).all()
index = jt.array([[0, 1, 2], [0, 1, 4]])
x = jt.zeros((3, 5), dtype=src.dtype).scatter_(1, index, src)
assert (x.data == [[1, 2, 3, 0, 0], [6, 7, 0, 0, 8], [0, 0, 0, 0,
0]]).all()
x = jt.full((2, 4), 2.).scatter_(1,
jt.array([[2], [3]]),
jt.array(1.23),
reduce='multiply')
assert np.allclose(x.data, [[2.0000, 2.0000, 2.4600, 2.0000],
[2.0000, 2.0000, 2.0000, 2.4600]]), x
x = jt.full((2, 4), 2.).scatter_(1,
jt.array([[2], [3]]),
jt.array(1.23),
reduce='add')
assert np.allclose(x.data, [[2.0000, 2.0000, 3.2300, 2.0000],
[2.0000, 2.0000, 2.0000, 3.2300]])
def vertex_normals(self):
if self._vertex_normals_update:
bs, nv = self.vertices.shape[:2]
bs, nf = self.faces.shape[:2]
faces = self.faces + (jt.arange(bs) * nv)[:, None, None]
vertices_faces = self.vertices.reshape((bs * nv, 3))[faces.long()]
faces = faces.view(-1, 3)
vertices_faces = vertices_faces.view(-1, 3, 3)
normals = (jt.cross(
vertices_faces[:, 2] - vertices_faces[:, 1],
vertices_faces[:, 0] - vertices_faces[:, 1])).reindex_reduce(
op="sum",
shape=[bs * nv, 3],
indexes=["@e0(i0)", "i1"],
extras=[faces[:, 1].long()]) + (jt.cross(
vertices_faces[:, 0] - vertices_faces[:, 2],
vertices_faces[:, 1] -
vertices_faces[:, 2])).reindex_reduce(
op="sum",
shape=[bs * nv, 3],
indexes=["@e0(i0)", "i1"],
extras=[faces[:, 2].long()]) + (jt.cross(
vertices_faces[:, 1] - vertices_faces[:, 0],
vertices_faces[:, 2] -
vertices_faces[:, 0])).reindex_reduce(
op="sum",
shape=[bs * nv, 3],
indexes=["@e0(i0)", "i1"],
extras=[faces[:, 0].long()])
normals = jt.normalize(normals, p=2, eps=1e-6, dim=1)
self._vertex_normals = normals.reshape((bs, nv, 3))
self._vertex_normals_update = False
return self._vertex_normals
def sample_images(batches_done):
"""Saves a generated sample from the validation set"""
batches = next(iter(val_dataloader))
real_A = batches[0]
real_A_eyel = batches[1]
real_A_eyer = batches[2]
real_A_nose = batches[3]
real_A_mouth = batches[4]
real_A_hair = batches[5]
real_A_bg = batches[6]
mask = batches[7]
mask2 = batches[8]
center = batches[9]
cmaskel = batches[10]
cmasker = batches[11]
cmaskno = batches[12]
cmaskmo = batches[13]
maskh = mask * mask2
maskb = inverse_mask(mask2)
fake_B0 = G_global(real_A)
# EYES, NOSE, MOUTH
fake_B_eyel1 = G_l_eyel(real_A_eyel)
fake_B_eyel2 = ae_eyel(fake_B_eyel1)
fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel)
fake_B_eyer1 = G_l_eyer(real_A_eyer)
fake_B_eyer2 = ae_eyer(fake_B_eyer1)
fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker)
fake_B_nose1 = G_l_nose(real_A_nose)
fake_B_nose2 = ae_nose(fake_B_nose1)
fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno)
fake_B_mouth1 = G_l_mouth(real_A_mouth)
outputs1 = CLm(real_A_mouth)
pred = jt.argmax(outputs1, dim=1)[0]
fake_B_mouth2w = ae_mowhite(fake_B_mouth1)
fake_B_mouth2b = ae_moblack(fake_B_mouth1)
fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1)
bs, c, h, w = fake_B_mouth2s.shape
index = pred + jt.arange(bs) * c
fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h,
w])[index].reshape([bs, 1, h, w])
fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo)
# HAIR & BG
outputs2 = CLh(real_A_hair)
onehot2 = getonehot(outputs2, 3, bs)
fake_B_hair = G_l_hair(real_A_hair, onehot2)
fake_B_bg = G_l_bg(real_A_bg)
# PARTCOMBINE
fake_B1 = partCombiner2_bg(center,
fake_B_eyel,
fake_B_eyer,
fake_B_nose,
fake_B_mouth,
fake_B_hair,
fake_B_bg,
maskh,
maskb,
comb_op=1,
load_h=opt.img_height,
load_w=opt.img_width)
# FUSION NET
fake_B = G_combine(jt.contrib.concat((fake_B0, fake_B1), 1))
img_sample = np.concatenate(
[real_A.data, fake_B.repeat(1, 3, 1, 1).data], -2)
save_image(img_sample,
"images/%s/%s.jpg" % (opt.dataset_name, batches_done),
nrow=5)
fake_B_eyel2 = ae_eyel(fake_B_eyel1)
fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel)
fake_B_eyer1 = G_l_eyer(real_A_eyer)
fake_B_eyer2 = ae_eyer(fake_B_eyer1)
fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker)
fake_B_nose1 = G_l_nose(real_A_nose)
fake_B_nose2 = ae_nose(fake_B_nose1)
fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno)
fake_B_mouth1 = G_l_mouth(real_A_mouth)
outputs1 = CLm(real_A_mouth)
pred = jt.argmax(outputs1, dim=1)[0]
fake_B_mouth2w = ae_mowhite(fake_B_mouth1)
fake_B_mouth2b = ae_moblack(fake_B_mouth1)
fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1)
bs, c, h, w = fake_B_mouth2s.shape
index = pred + jt.arange(bs) * c
fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h, w])[index].reshape(
[bs, 1, h, w])
fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo)
# HAIR & BG
outputs2 = CLh(real_A_hair)
onehot2 = getonehot(outputs2, 3, bs)
fake_B_hair = G_l_hair(real_A_hair, onehot2)
fake_B_bg = G_l_bg(real_A_bg)
# PARTCOMBINE
fake_B1 = partCombiner2_bg(center,
fake_B_eyel,
fake_B_eyer,
fake_B_nose,
fake_B_mouth,
fake_B_hair,
def read_planetoid_data(folder, prefix):
names = ['x', 'tx', 'allx', 'y', 'ty', 'ally', 'graph', 'test.index']
items = [read_file(folder, prefix, name) for name in names]
x, tx, allx, y, ty, ally, graph, test_index = items
train_index = jt.arange(y.size(0), dtype=Var.int32)
val_index = jt.arange(y.size(0), y.size(0) + 500, dtype=Var.int32)
sorted_test_index = test_index.argsort()[1]
if prefix.lower() == 'citeseer':
# There are some isolated nodes in the Citeseer graph, resulting in
# none consecutive test indices. We need to identify them and add them
# as zero vectors to `tx` and `ty`.
len_test_indices = (test_index.max() - test_index.min()).item() + 1
tx_ext = jt.zeros((len_test_indices, tx.size(1)))
tx_ext[sorted_test_index - test_index.min(), :] = tx
ty_ext = jt.zeros((len_test_indices, ty.size(1)))
ty_ext[sorted_test_index - test_index.min(), :] = ty
tx, ty = tx_ext, ty_ext
if prefix.lower() == 'nell.0.001':
tx_ext = jt.zeros((len(graph) - allx.size(0), x.size(1)))
tx_ext[sorted_test_index - allx.size(0)] = tx
ty_ext = jt.zeros((len(graph) - ally.size(0), y.size(1)))
ty_ext[sorted_test_index - ally.size(0)] = ty
tx, ty = tx_ext, ty_ext
x = jt.concat([allx, tx], dim=0)
x[test_index] = x[sorted_test_index]
# Creating feature vectors for relations.
row, col, value = SparseTensor.from_dense(x).coo()
rows, cols, values = [row], [col], [value]
mask1 = index_to_mask(test_index, size=len(graph))
mask2 = index_to_mask(jt.arange(allx.size(0), len(graph)),
size=len(graph))
mask = jt.logical_or(jt.logical_not(mask1), jt.logical_not(mask2))
isolated_index = mask.nonzero(as_tuple=False).view(-1)[allx.size(0):]
rows += [isolated_index]
cols += [jt.arange(isolated_index.size(0)) + x.size(1)]
values += [jt.ones((isolated_index.size(0)))]
x = SparseTensor(row=jt.concat(rows),
col=jt.concat(cols),
value=jt.concat(values))
else:
x = jt.concat([allx, tx], dim=0)
x[test_index] = x[sorted_test_index]
y = jt.concat([ally, ty], dim=0).argmax(dim=1)[0]
y[test_index] = y[sorted_test_index]
train_mask = index_to_mask(train_index, size=y.size(0))
val_mask = index_to_mask(val_index, size=y.size(0))
test_mask = index_to_mask(test_index, size=y.size(0))
edge_index = edge_index_from_dict(graph, num_nodes=y.size(0))
data = Data(x=x, edge_index=edge_index, y=y)
data.train_mask = train_mask
data.val_mask = val_mask
data.test_mask = test_mask
return data
def __init__(self, n, weight=False): # n: number of inputs
super(Sum, self).__init__()
self.weight = weight # apply weights boolean
self.iter = range(n - 1) # iter object
if weight:
self.w = -jt.arange(1., n) / 2 # layer weights
def test_ellipsis_with_none(self):
a = jt.arange(2 * 4 * 4).reshape(2, 4, 4)
b = a[..., :, None, :2]
assert b.shape == [2, 4, 1, 2]
np.testing.assert_allclose(b.data, a.data[..., :, None, :2])
def forward_for_single_feature_map(self, anchors, objectness,
box_regression):
"""
Arguments:
anchors: list[BoxList]
objectness: tensor of size N, A, H, W
box_regression: tensor of size N, A * 4, H, W
"""
# global II
# import pickle
N, A, H, W = objectness.shape
# put in the same format as anchors
objectness = permute_and_flatten(objectness, N, A, 1, H,
W).reshape(N, -1)
# print('objectness',objectness.mean())
objectness = objectness.sigmoid()
box_regression = permute_and_flatten(box_regression, N, A, 4, H, W)
# print('regression',box_regression.mean())
num_anchors = A * H * W
pre_nms_top_n = min(self.pre_nms_top_n, num_anchors)
# print(pre_nms_top_n)
#print('objectness',objectness)
# objectness = jt.array(pickle.load(open(f'/home/lxl/objectness_0_{II}.pkl','rb')))
# print(objectness.shape)
objectness, topk_idx = objectness.topk(pre_nms_top_n,
dim=1,
sorted=True)
# print(II,'topk',topk_idx.sum(),topk_idx.shape)
batch_idx = jt.arange(N).unsqueeze(1)
# pickle.dump(topk_idx.numpy(),open(f'/home/lxl/topk_idx_{II}_jt.pkl','wb'))
# topk_idx_tmp = topk_idx.numpy()
# batch_idx = jt.array(pickle.load(open(f'/home/lxl/batch_idx_{II}.pkl','rb')))
# topk_idx = jt.array(pickle.load(open(f'/home/lxl/topk_idx_{II}.pkl','rb')))
# err = np.abs(topk_idx_tmp-topk_idx.numpy())
# print('Error!!!!!!!!!!!!!!!!',err.sum())
# print(err.nonzero())
#print('box_regression0',box_regression)
#batch_idx = jt.index(topk_idx.shape,dim=0)
box_regression = box_regression[batch_idx, topk_idx]
#print('box_regression1',box_regression)
image_shapes = [box.size for box in anchors]
concat_anchors = jt.contrib.concat([a.bbox for a in anchors], dim=0)
concat_anchors = concat_anchors.reshape(N, -1, 4)[batch_idx, topk_idx]
# box_regression = jt.array(pickle.load(open(f'/home/lxl/box_regression_{II}.pkl','rb')))
# concat_anchors = jt.array(pickle.load(open(f'/home/lxl/concat_anchors_{II}.pkl','rb')))
proposals = self.box_coder.decode(box_regression.reshape(-1, 4),
concat_anchors.reshape(-1, 4))
proposals = proposals.reshape(N, -1, 4)
# proposals = jt.array(pickle.load(open(f'/home/lxl/proposal_{II}.pkl','rb')))
# objectness = jt.array(pickle.load(open(f'/home/lxl/objectness_{II}.pkl','rb')))
# II+=1
result = []
for i in range(len(image_shapes)):
proposal, score, im_shape = proposals[i], objectness[
i], image_shapes[i]
boxlist = BoxList(proposal, im_shape, mode="xyxy")
boxlist.add_field("objectness", score)
boxlist = boxlist.clip_to_image(remove_empty=False)
boxlist = remove_small_boxes(boxlist, self.min_size)
boxlist = boxlist_nms(
boxlist,
self.nms_thresh,
max_proposals=self.post_nms_top_n,
score_field="objectness",
)
result.append(boxlist)
return result
