def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
if args.norm:
transforms.append(
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
return T.Compose(transforms)
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
# get the number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
# Test transform
test_transform = transforms.Compose([
transforms.ToTensor()
])
# create dataset and set seed to recreate results
dataset = TemplateDataset("dataset", test_transform)
torch.manual_seed(1)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# move model to the right device
model.to(device)
# pick one image from the test set
img, _ = dataset[0]
# put the model in evaluation mode
# model = torch.nn.DataParallel(model)
type=int,
default=300,
help='used during inference')
parser.add_argument('--resume', default='', help='resume from checkpoint')
args = parser.parse_args()
exp_name = (args.resume).split("/")[-2]
out_dir = os.path.join('../jsons', args.dataset, exp_name)
utils.mkdir(out_dir)
root = '../../../../datasets/coco/images'
if args.dataset == 'lvis':
annotations = "../../../../datasets/coco/annotations/lvis_v1_val.json"
dset = LVISDetection(root,
annotations,
transforms=transforms.ToTensor())
num_classes = 1204
elif args.dataset == 'coco':
annotations = "../../../../datasets/coco/annotations/instances_val2017.json"
dset = CocoDetection(root,
annotations,
transforms=transforms.ToTensor())
num_classes = 91
else:
sys.exit("Dataset not recognisable")
if args.model == 'fasterrcnn_resnet50_fpn':
model = frcnn.fasterrcnn_resnet50_fpn(
pretrained=False,
num_classes=num_classes,
box_score_thresh=float(args.confidence_threshold),