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"))