def main():
# Load configuration
args = parser.parse_args()
# Torch stuff
torch.cuda.set_device(args.rank)
cudnn.benchmark = True
# Create model by loading a snapshot
body, head, cls_state = load_snapshot(args.snapshot)
model = SegmentationModule(body, head, 256, 65, args.fusion_mode)
model.cls.load_state_dict(cls_state)
model = model.cuda().eval()
print(model)
# Create data loader
transformation = SegmentationTransform(
2048,
(0.41738699, 0.45732192, 0.46886091),
(0.25685097, 0.26509955, 0.29067996),
)
dataset = SegmentationDataset(args.data, transformation)
data_loader = DataLoader(
dataset,
batch_size=1,
pin_memory=True,
sampler=DistributedSampler(dataset, args.world_size, args.rank),
num_workers=2,
collate_fn=segmentation_collate,
shuffle=False,
)
# Run testing
scales = eval(args.scales)
with torch.no_grad():
for batch_i, rec in enumerate(data_loader):
print("Testing batch [{:3d}/{:3d}]".format(batch_i + 1,
len(data_loader)))
img = rec["img"].cuda(non_blocking=True)
probs, preds = model(img, scales, args.flip)
for i, (prob, pred) in enumerate(
zip(torch.unbind(probs, dim=0), torch.unbind(preds,
dim=0))):
out_size = rec["meta"][i]["size"]
img_name = rec["meta"][i]["idx"]
# Save prediction
prob = prob.cpu()
pred = pred.cpu()
pred_img = get_pred_image(pred, out_size,
args.output_mode == "palette")
pred_img.save(path.join(args.output, img_name + ".png"))
# Optionally save probabilities
if args.output_mode == "prob":
prob_img = get_prob_image(prob, out_size)
prob_img.save(
path.join(args.output, img_name + "_prob.png"))
def main():
#object_class = [20,21,52,53,54,55,56,57,58,59,60,61,62,64]
# Load configuration
args = parser.parse_args()
# Torch stuff
torch.cuda.set_device(args.rank)
cudnn.benchmark = True
# Create model by loading a snapshot
body, head, cls_state = load_snapshot(args.snapshot)
model = SegmentationModule(body, head, 256, 65, args.fusion_mode)
model.cls.load_state_dict(cls_state)
model = model.cuda().eval()
# Create data loader
transformation = SegmentationTransform(
1920,
(0.41738699, 0.45732192, 0.46886091),
(0.25685097, 0.26509955, 0.29067996),
)
for dir_ in os.listdir('/data/JAAD_clip_images'):
dir_ = dir_.split('.')[0]
def main():
#object_class = [20,21,52,53,54,55,56,57,58,59,60,61,62,64]
# Load configuration
args = parser.parse_args()
# Torch stuff
torch.cuda.set_device(args.rank)
cudnn.benchmark = True
# Create model by loading a snapshot
body, head, cls_state = load_snapshot(args.snapshot)
model = SegmentationModule(body, head, 256, 65, args.fusion_mode)
model.cls.load_state_dict(cls_state)
model = model.cuda().eval()
# Create data loader
transformation = SegmentationTransform(
1920,
(0.41738699, 0.45732192, 0.46886091),
(0.25685097, 0.26509955, 0.29067996),
)
for dir_ in os.listdir('/data/JAAD/JAAD_clip_images'):
dir_ = dir_.split('.')[0]
os.system('mkdir -p /workspace/caozhangjie/inplace_abn/JAAD_features_occl/'+dir_)
os.system('mkdir -p /workspace/caozhangjie/inplace_abn/JAAD_pos_occl/'+dir_)
os.system('mkdir -p /workspace/caozhangjie/inplace_abn/JAAD_seg_occl/'+dir_)
dataset = SegmentationDataset('/data/JAAD/JAAD_clip_images', '/workspace/caozhangjie/inplace_abn/JAAD_vbb/vbb_full', '/workspace/caozhangjie/inplace_abn/JAAD_seg_occl', '/workspace/caozhangjie/inplace_abn/JAAD_features_occl', '/workspace/caozhangjie/inplace_abn/JAAD_pos_occl', transformation, "train", [args.start,args.start+10], 100)
data_loader = DataLoader(
dataset,
batch_size=1,
pin_memory=True,
num_workers=4,
collate_fn=segmentation_collate,
shuffle=False
)
# Run testing
scales = [int(val) for val in args.scales.split(",")]
with torch.no_grad():
for batch_i, rec in enumerate(data_loader):
print("Testing batch [{:3d}/{:3d}]".format(batch_i + 1, len(data_loader)))
img = rec["img"].cuda(non_blocking=True)
x_up, (probs, preds) = model(img, scales, args.flip)
x_up = x_up.cpu()
for i, (prob, pred) in enumerate(zip(torch.unbind(probs, dim=0), torch.unbind(preds, dim=0))):
out_size = rec["size"][i]
bboxes = rec["pos"][i]
ped_appear = rec['ped_appear'][i]
#img_name = rec["meta"][i]["idx"]
# Save prediction
prob = prob.cpu()
pred = pred.cpu()
pred_img = get_pred_image(pred, out_size, args.output_mode == "palette")
pred_img.save(rec["seg_out"][i])
#torch.save(x_up.cpu(), rec['raw_feature_out'][i])
x_up = x_up.cpu()
width_dict, height_dict = get_box(pred)
features = []
spp_scales = [4,2,1]
height_threshold = 15
width_threshold = 15
feature_dim = x_up.size(1) * sum([scale * scale for scale in spp_scales])
position = []
for j in range(65):
if j != 19:
if width_dict[j][0] != -1:
if not ((width_dict[j][1]+1 - width_dict[j][0] < width_threshold) or (height_dict[j][1]+1 - height_dict[j][0] < height_threshold)):
raw_feature = (x_up * ((pred == j).float()).view(1, 1, pred.size(-2), pred.size(-1)))[:,:,height_dict[j][0]:height_dict[j][1]+1, width_dict[j][0]:width_dict[j][1]+1]
position.append([width_dict[j][0], height_dict[j][0], \
width_dict[j][1]+1-width_dict[j][0], height_dict[j][1]+1-height_dict[j][0]])
spp_feature = spatial_pyramid_pool(raw_feature.cuda().detach(), 1, [raw_feature.size(2), raw_feature.size(3)], [4, 2, 1]).cpu()
features.append(spp_feature)
else:
features.append(torch.zeros([1, feature_dim]))
position.append([-1, -1, -1, -1])
else:
features.append(torch.zeros([1, feature_dim]))
position.append([-1, -1, -1, -1])
for j in range(ped_appear.shape[0]):
bbox = bboxes[j, :]
if ped_appear[j] > 0 and (bbox[2]>=width_threshold and bbox[3] >=height_threshold):
ped_raw_feature = x_up[:,:,bbox[1]:bbox[1]+bbox[3], \
bbox[0]:bbox[0]+bbox[2]]
features.append(spatial_pyramid_pool(ped_raw_feature.cuda().detach(), \
1, [ped_raw_feature.size(2), ped_raw_feature.size(3)], [4, 2, 1]).cpu())
else:
features.append(torch.zeros([1, feature_dim]))
position.append([bbox[i] for i in range(4)])
features = torch.cat(features, dim=0)
position = torch.from_numpy(np.array(position))
#print(features.size())
#print(position.size())
torch.save(features, rec['feature_out'][i])
torch.save(position, rec['pos_out'][i])
def main():
# Load configuration
args = parser.parse_args()
# Torch stuff
device = torch.device("cuda:0")
cudnn.benchmark = True
# Create model by loading a snapshot
body, head, cls_state = load_snapshot(args.snapshot)
model = SegmentationModule(body, head, 256, 65, args.fusion_mode)
model.cls.load_state_dict(cls_state)
model = nn.DataParallel(model, output_device=-1).cuda(device).eval()
print(model)
# Create data loader
transformation = SegmentationTransform(
2048,
(1024, 2048),
(0.41738699, 0.45732192, 0.46886091),
(0.25685097, 0.26509955, 0.29067996),
)
dataset = SegmentationDataset(args.data, transformation)
data_loader = DataLoader(dataset,
batch_size=torch.cuda.device_count(),
pin_memory=True,
num_workers=torch.cuda.device_count(),
collate_fn=segmentation_collate,
shuffle=False)
# Run testing
scales = eval(args.scales)
with torch.no_grad():
for batch_i, rec in enumerate(data_loader):
print("Testing batch [{:3d}/{:3d}]".format(batch_i + 1,
len(data_loader)))
img = rec["img"].cuda(device, True)
probs, preds = model(img, scales, args.flip)
for i, (prob, pred) in enumerate(
zip(torch.unbind(probs, dim=0), torch.unbind(preds,
dim=0))):
crop_bbx = rec["meta"][i]["valid_bbx"]
out_size = rec["meta"][i]["size"]
img_name = rec["meta"][i]["idx"]
# Crop to valid area
pred = pred[crop_bbx[0]:crop_bbx[2], crop_bbx[1]:crop_bbx[3]]
# Save prediction
pred_img = get_pred_image(pred, out_size,
args.output_mode == "palette")
pred_img.save(path.join(args.output, img_name + ".png"))
# Optionally save probabilities
if args.output_mode == "prob":
prob = prob[crop_bbx[0]:crop_bbx[2],
crop_bbx[1]:crop_bbx[3]]
prob_img = get_prob_image(prob, out_size)
prob_img.save(
path.join(args.output, img_name + "_prob.png"))
