class SegmentationNetwork(nn.Module):
def __init__(self,
backbone=None,
output_channels=2,
backbone_output_channels=512):
super().__init__()
if not backbone:
b, _ = remove_backbone_head(resnet18(pretrained=False))
backbone = b
# ResNet produces 512-length outputs
backbone.out_channels = backbone_output_channels
self.segmentation_network = FasterRCNN(backbone,
num_classes=output_channels)
def forward(self, x, boxes=None):
x = self.segmentation_network(x, boxes)
return x
def infer(self, x):
self.eval()
self.segmentation_network.eval()
x = self.segmentation_network(x)
return convert_bounding_box_inference(x)
def fasterrcnn_resnetxx_fpnxx(cfg):
backbone = resnet.__dict__[cfg['backbone_name']](
pretrained=cfg['backbone_pretrained'],
norm_layer=misc_nn_ops.FrozenBatchNorm2d)
# freeze layers
for name, parameter in backbone.named_parameters():
if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
parameter.requires_grad_(False)
# wrapper backbone with fpn
return_layers = cfg['fpn']['return_layers']
in_channels_stage2 = backbone.inplanes // 8
in_channels_list = [in_channels_stage2 * 2**i for i in range(len(return_layers))]
out_channels = cfg['fpn']['out_channels']
backbone_fpn = BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels)
anchor_generator = AnchorGenerator(**cfg['anchor_generator'])
# print(anchor_generator.num_anchors_per_location())
roi_pooler = MultiScaleRoIAlign(**cfg['box_roi_pool'])
model = FasterRCNN(backbone_fpn, num_classes=cfg['num_classes'], rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
if os.path.exists(cfg['fasterrcnn_pretrained']):
state_dict = torch.load(cfg['fasterrcnn_pretrained'])
model.load_state_dict(state_dict)
return model
def __init__(self, h_g, h_l, g, k, s, c):
super(glimpse_network, self).__init__()
self.retina = retina(g, k, s)
self.feature_extractor = nn.Sequential(
*list(resnet50(pretrained=True).children())[:-1])
#detection feature
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
self.detection_model = FasterRCNN(
backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.detection_model.eval()
# glimpse layer
D_in = k * g * g * c
D_in = 2048
self.fc1 = nn.Linear(D_in, h_g)
# location layer
D_in = 2
self.fc2 = nn.Linear(D_in, h_l)
self.fc3 = nn.Linear(h_g, h_g + h_l)
self.fc4 = nn.Linear(h_l, h_g + h_l)
def get_model(
backbone_name="resnet50",
detector_name="fasterrcnn",
trainable_layers=3,
model_ckpt=None,
):
"""Constructs a fasterrcnn or maskrcnn detector with the given backbone"""
num_classes = 2 # 1 class (wheat) + background
if model_ckpt:
# backbone = resnet_fpn_backbone('resnet101', True)
backbone = timm_resnet_fpn_backbone(backbone_name, False,
trainable_layers)
else:
backbone = timm_resnet_fpn_backbone(backbone_name, True,
trainable_layers)
if detector_name == "fasterrcnn":
model = FasterRCNN(backbone, num_classes)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
elif detector_name == "maskrcnn":
model = MaskRCNN(backbone, num_classes)
in_features_mask = (
model.roi_heads.mask_predictor.conv5_mask.in_channels)
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, num_classes)
else:
raise Exception(f"{detector_name} is not supported")
if model_ckpt is not None:
model.load_state_dict(torch.load(model_ckpt)["model_state_dict"])
print("loaded ckpt")
return model
def start_evaluation(test_data_loader: DataLoader, model: FasterRCNN,
device: str, epoch: int, logger: Logger,
args: argparse.Namespace) -> Tensor:
"""
Evaluate the model with the test set
:param test_data_loader: Data loader for test data:
:param model: Model that is being tested
:param device: Device for the computation
:param epoch: Current epoch
:param logger: Logger for logging handling
:param args: Arguments
:return:
"""
logger.info(f'Start evaluation after {epoch} epochs')
model.eval()
scores = []
for idx, result in enumerate(test_data_loader):
images = list(image.to(device) for image in result[0])
targets = result[1]
with torch.set_grad_enabled(False):
outputs = model(images)
for output_idx, element in enumerate(outputs):
predicted_labels = element['labels']
true_labels = targets[output_idx]['labels']
if len(element['scores']) != 0:
scores.append(torch.mean(element['scores']))
if idx % args.print_status:
logger.info(
f'Scores {element["scores"]} \n'
f'Labels predicted: {predicted_labels} Groundtruth labels: {true_labels}'
)
avg_score = torch.mean(torch.Tensor(scores))
return avg_score
def train_one_epoch(train_data_loader: DataLoader, model: FasterRCNN,
device: str, logger: Logger, optimizer: SGD,
args: argparse.Namespace) -> Tuple[float, float]:
"""
Training the model for one epoch
:param train_data_loader: Dataloader for training data
:param model: FasterRCNN model
:param device: Device for the computation
:param logger: Logger for logging handling
:param optimizer: Optmizer for the training
:param args: Arguments
:return:
"""
start = time.time()
model.train()
loss_per_iteration = []
for idx, result in enumerate(train_data_loader):
images = list(image.to(device) for image in result[0])
targets = result[1]
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
optimizer.zero_grad()
with torch.set_grad_enabled(True):
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
losses.backward()
optimizer.step()
loss_per_iteration.append(losses)
loss_doc_str = "".join("{}:{} ".format(key, val)
for key, val in loss_dict.items())
if idx % args.print_status == 0:
logger.info(f'Iteration: [{idx}/{len(train_data_loader)}]\t'
f'Loss: {losses} \t'
f'Loss_dict: {loss_doc_str}')
epoch_time = time.time() - start
return loss_per_iteration, epoch_time
def demo():
# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
backbone.out_channels = 1280
# let's make the RPN generate 5 x 3 anchors per spatial
# location, with 5 different sizes and 3 different aspect
# ratios. We have a Tuple[Tuple[int]] because each feature
# map could potentially have different sizes and
# aspect ratios
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)
print(predictions)
class ResNet50_FasterRCNN:
def __init__(self, pretrained=False):
# Building our FasterRCNN model for objects detection
backbone = resnet_fpn_backbone('resnet50', pretrained=pretrained)
num_classes = 4 + 1
anchor_generator = AnchorGenerator(sizes=(40, 60, 150, 200, 250),
aspect_ratios=(0.7, 1.0, 1.3))
self.model = FRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator)
def train(self):
self.model.train()
def to(self, device):
self.model.to(device)
def eval(self):
self.model.eval()
def parameters(self):
return self.model.parameters()
def get_state_dict(self):
return self.model.state_dict()
def set_state_dict(self, state_dict):
self.model.load_state_dict(state_dict)
def fit_batch(self, images, target):
return self.model(images, target)
def predict_batch(self, images):
return self.model(images)
def __init__(self,
backbone_name: str,
pretrained: bool = True,
finetune: bool = True,
num_classes: int = 2):
self.__pretrained = pretrained
self.__num_classes = num_classes
self.__model_name = backbone_name
backbone = build_backbone(backbone_name, pretrained, finetune)
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
self.model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.params = [p for p in self.model.parameters() if p.requires_grad]
self.optimizer = torch.optim.Adam(params=self.params,
lr=0.005,
weight_decay=0.0005)
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=self.optimizer, step_size=3, gamma=0.1)
def train_mask():
torch.manual_seed(1)
dataset = unimib_data.UNIMIBDataset(get_transform(train=True))
dataset_test = unimib_data.UNIMIBDataset(get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=1, shuffle=True,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False,
collate_fn=utils.collate_fn)
# device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
device = "cpu"
num_classes = 74
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
for param in backbone:
param.requires_grad = False
backbone.out_channels = 1280
model.backbone = backbone
model = FasterRCNN(backbone,
num_classes=74)
model = torch.nn.Sequential(
model.rpn,
model.roi_heads
)
print(model)
# for param in model.backbone:
# param.requires_grad = False
# model = get_instance_segmentation_model(num_classes)
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
momentum=0.9, weight_decay=0.0005)
# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 1
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
# train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# update the learning rate
# lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
def __init__(self, config, logger):
self._logger = logger
self._threshold = config['threshold']
modelfile = config['model']
self._device = config['device'] # cpu, cuda, cuda:0
backbone = resnet_fpn_backbone('resnet50', False)
self._model = FasterRCNN(backbone, 8) # 8 classes
checkpoint = torch.load(modelfile, map_location=self._device)
self._model.load_state_dict(checkpoint['model_state_dict'])
device = torch.device(self._device)
self._model.to(device)
self._model.eval()
def fasterrcnn_resnet_fpn(backbone_name,
pretrained=False,
trainable_backbone_layers=None,
num_classes=81,
pretrained_backbone=True,
**kwargs):
"""
Construct Faster R-CNN with a ResNet-FPN backbone
Arguments:
backbone_name(str): Name of the backbone.
Refer to torchvision.models.resnet.__dict__ for details
pretrained(bool, optional): If True, load weights for the detector
pretrained on MS COCO. Only ResNet50-FPN is supported for the moment.
trainable_backbone_layers(int, optional): Number of trainable (not frozen)
resnet layers starting from final block.
num_classes(int, optional): Number of target classes.
pretrained_backbone(bool, optional): If True, load weights for backbone
pre-trained on ImageNet
Refer to torchvision.models.detection.FasterRCNN for kwargs
"""
trainable_backbone_layers = _validate_trainable_layers(
pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3)
if pretrained and backbone_name == 'resnet50':
# no need to download the backbone if pretrained is set
pretrained_backbone = False
backbone = resnet_fpn_backbone(backbone_name, pretrained_backbone)
model = FasterRCNN(backbone, num_classes, **kwargs)
if pretrained and backbone_name == 'resnet50':
state_dict = models.utils.load_state_dict_from_url(
model_urls['fasterrcnn_resnet50_fpn_coco'])
if num_classes == 81:
# Remove the parameters for the additional classes
state_dict['roi_heads.box_predictor.cls_score.weight'] = \
state_dict['roi_heads.box_predictor.cls_score.weight'][KEEP]
state_dict['roi_heads.box_predictor.cls_score.bias'] = \
state_dict['roi_heads.box_predictor.cls_score.bias'][KEEP]
state_dict['roi_heads.box_predictor.bbox_pred.weight'] = \
state_dict['roi_heads.box_predictor.bbox_pred.weight'][KEEPX4.flatten()]
state_dict['roi_heads.box_predictor.bbox_pred.bias'] = \
state_dict['roi_heads.box_predictor.bbox_pred.bias'][KEEPX4.flatten()]
model.load_state_dict(state_dict)
elif pretrained:
print(
"WARNING: No pretrained detector on MS COCO with {}.".format(
backbone_name),
"Proceed with only pretrained backbone on ImageNet.")
return model
def __init__(self,
backbone=None,
output_channels=2,
backbone_output_channels=512):
super().__init__()
if not backbone:
b, _ = remove_backbone_head(resnet18(pretrained=False))
backbone = b
# ResNet produces 512-length outputs
backbone.out_channels = backbone_output_channels
self.segmentation_network = FasterRCNN(backbone,
num_classes=output_channels)
def get_model_detection(num_classes,
model_name="resnet101",
aspect_ratios=(0.5, 1.0, 2.0),
anchor_sizes=((32, ), (64, ), (128, ), (256, ),
(512, )),
backbone_chkpt=None,
model_chkpt=None,
train_backbone=True,
**kwargs):
"""
"""
backbone = load_backbone(model_name, backbone_chkpt)
anchors = get_anchors(aspect_ratios=aspect_ratios,
anchor_sizes=anchor_sizes)
model = FasterRCNN(backbone,
num_classes,
rpn_anchor_generator=anchors,
**kwargs)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
if model_chkpt:
backbone.load_state_dict(torch.load(model_chkpt, map_location="cpu"))
if not train_backbone:
for p in model.backbone.parameters():
p.requires_grad_(False)
return model
def Model(self, model_name, use_pretrained=True, use_gpu=True):
self.system_dict["model_name"] = model_name
self.system_dict["use_pretrained"] = use_pretrained
self.system_dict["use_gpu"] = use_gpu
if (self.system_dict["model_name"] in self.system_dict["model_set_1"]):
first_name, second_name = self.system_dict["model_name"].split("_")
if (first_name == "faster-rcnn" and second_name == "mobilenet-v2"):
backbone = torchvision.models.mobilenet_v2(
pretrained=use_pretrained).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256,
512), ),
aspect_ratios=((0.5, 1.0,
2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=[0], output_size=7, sampling_ratio=2)
self.system_dict["local"]["model"] = FasterRCNN(
backbone,
num_classes=self.system_dict["local"]["num_classes"],
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.set_device(use_gpu=self.system_dict["use_gpu"])
self.system_dict["local"]["model"].to(
self.system_dict["local"]["device"])
def FasterRCNN_resnext50_32x4d():
net = torchvision.models.resnext50_32x4d(pretrained=True)
modules = list(net.children())[:-1]
backbone = nn.Sequential(*modules)
# backbone = torchvision.models.resnext50_32x4d(pretrained=True).features
# test_backbone(backbone)
backbone.out_channels = 2048
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=config.CLASSES,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def fasterrcnn_mobilenetv2_fpn(pretrained=False,
progress=True,
num_classes=91,
pretrained_backbone=True,
**kwargs):
"""
Construct a Faster RCNN model with mobilenetv2 backbone
"""
if pretrained:
pretrained_backbone = False
backbone = torchvision.models.mobilenet_v2(
pretrained=pretrained_backbone).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def fasterrcnn_resnet34(num_classes, **kwargs):
backbone = resnet_fpn_backbone("resnet34", pretrained=True)
model = FasterRCNN(backbone, num_classes, **kwargs)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
def get_model_instance_segmentation2(num_classes):
# COCO 에서 미리 학습된 인스턴스 분할 모델을 읽어옵니다
#model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
#model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
backbone = torchvision.models.mobilenet_v2(pretrained=False).features
#backbone.out_channels = 1
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=1,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
print("mobilenet_v2 call2 - out_channels :1280, 19,540,921")
# 분류를 위한 입력 특징 차원을 얻습니다
#in_features = backbone
# 미리 학습된 헤더를 새로운 것으로 바꿉니다
#model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
#in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
#hidden_layer = 1
# and replace the mask predictor with a new one
#model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
# hidden_layer,
# num_classes)
return model
def model_predictions_from_loader(
model: FasterRCNN,
data_loader: DataLoader,
device: torch.device,
max_n_batches=1,
normalization_transform: Normalize = Normalize([1, 1, 1], [1, 1, 1])
) -> List[Tuple[BoxedExample, BoxedExample]]:
model = model.to(device).eval()
input_examples = []
output_examples = []
with torch.no_grad():
for i, (img, labels) in enumerate(data_loader):
if i >= max_n_batches:
break
outputs = model([i.to(device) for i in img])
input_exs = [
model_input_to_boxed_example((im, l), normalization_transform)
for (im, l) in zip(img, labels)
]
output_exs = [
i.replace(boxes=o["boxes"].cpu().numpy())
for (i, o) in zip(input_exs, outputs)
]
input_examples.extend(input_exs)
output_examples.extend(output_exs)
return list(zip(input_examples, output_examples))
def faster_rcnn_alt_backbone():
# 特徴量のみを取得
backbone = torchvision.models.mobilentet_v2(pretrained=True).features
# FasterRCNNはバックボーンからの出力チャネル数を知る必要がある
# mobilenet_v2の出力は1280
backbone.out_channels = 1280
# RPN: Resion Proposal Networkに空間ごとにアンカーを生成。
# 5つのサイズと3つのアスペクト比が存在することを意味します。
# 特徴マップごとに異なるサイズとアスペクト比となる可能性があるので，
# Tuple[Tuple[int]] という形式で指定します。
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# 関心領域のトリミングを実行するために使用する特徴マップ(featmap_names)と
# 画像の大きさをもとに戻したあとのトリミングのサイズ(output_size)を適宜しましょう。
# バックボーンがテンソルを返す場合、featmap_nameは[0]になっているはずです。
# もう少し一般化して説明すると、バックボーンはOrderedDict[Tensor]を返すことになるので、
# featmap_namesで使用する特徴マップを選択できます。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def __init__(self, pretrained_path):
super().__init__()
self.pretrained_path = pretrained_path
self.output_dim = 800 * 800
# ------------------
# PRE-TRAINED MODEL
# ------------------
ae = BasicAE.load_from_checkpoint(pretrained_path)
ae.freeze()
self.backbone = ae.encoder
self.backbone.c3_only = True
self.backbone.out_channels = 32
# ------------------
# FAST RCNN
# ------------------
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
self.fast_rcnn = FasterRCNN(self.backbone,
num_classes=9,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# for unfreezing encoder later
self.frozen = True
def __init__(self, cfg, **kwargs):
super(WheatDetector, self).__init__()
self.backbone = fpn_backbone(pretrained=True)
self.base = FasterRCNN(self.backbone,
num_classes=cfg.MODEL.NUM_CLASSES,
**kwargs)
self.base.roi_heads.fastrcnn_loss = self.fastrcnn_loss
def fasterrcnn_resnet(backbone_name,
num_classes=91,
pretrained_backbone=True,
**kwargs):
"""
Construct Faster R-CNN with a ResNet backbone
Arguments:
backbone_name(str): Name of the backbone.
Refer to torchvision.models.resnet.__dict__ for details
num_classes(int, optional): Number of target classes, default: 91(COCO)
pretrained_backbone(bool, optional): If True, load weights for backbone
pre-trained on ImageNet
Refer to torchvision.models.detection.FasterRCNN for kwargs
"""
backbone, res5 = resnet_backbone(backbone_name, pretrained_backbone)
backbone.out_channels = 1024
box_head = nn.Sequential(res5, nn.AdaptiveAvgPool2d((1, 1)))
box_predictor = FastRCNNPredictor(2048, num_classes)
model = FasterRCNN(backbone,
box_head=box_head,
box_predictor=box_predictor,
**kwargs)
return model
def create_model(num_classes=3):
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=["0"],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
# pretrained=False,
# pretrained_backbone=True,
# num_classes=num_classes
# )
# model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# # get 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)
return model
def __call__(self,
classes=3,
sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), )):
from torchvision.models.detection.rpn import AnchorGenerator
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.faster_rcnn import TwoMLPHead, FastRCNNPredictor
# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.squeezenet1_1(pretrained=True).features
# FasterRCNN needs to know the number of
# output channels in a backbone. For squeezenet1_1, it's 512
# so we need to add it here
backbone.out_channels = 512
anchor_generator = AnchorGenerator(sizes=sizes,
aspect_ratios=aspect_ratios)
roi_out_size = 7
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'], output_size=roi_out_size, sampling_ratio=2)
representation_size = 256 # Scaled down from 1024 in original implementation.
# allows to reduce considerably the number of parameters
box_head = TwoMLPHead(backbone.out_channels * roi_out_size**2,
representation_size)
box_predictor = FastRCNNPredictor(representation_size, classes)
model = FasterRCNN(backbone,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
box_head=box_head,
box_predictor=box_predictor)
return model
def __init__(self):
super(DetectionModel, self).__init__()
# Creating the feature extractor
self.feature_extractor = DigitModel()
self.feature_extractor.out_channels = 512
# The region proposal network will generate 5x3 anchors per spatial location
# Each with 5 different sizes and 3 different aspect ratios
self.anchor_generator = AnchorGenerator(sizes=((16, 32, 64, 128,
256), ),
aspect_ratios=((0.5, 1.0,
2.0), ))
# Extract and align features from each proposed region on each scale.
# These features would then be sent to the final fully connected layers for classification and bbox regression
self.roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# We define 11 classes, 1 for each digit 0-9 and an additional class for the background
self.num_classes = 11
# Compile these modules within a PyTorch implementation of Faster RCNN
self.detector = FasterRCNN(self.feature_extractor,
num_classes=self.num_classes,
rpn_anchor_generator=self.anchor_generator,
box_roi_pool=self.roi_pooler,
min_size=384,
max_size=512,
box_score_thresh=0.5,
box_nms_thresh=0.5,
box_detections_per_img=5)
def __init__(self, num_classes):
super(TrashLiteModel, self).__init__()
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
self.model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
def resnet(n_classes, backbone_n='resnet50', pretrained_backbone=True):
# FIX doesn't work
"""
Faster-RCNN with
:param n_classes: number of classes of Fast-RCNN-Predictor
:param backbone_n: name of backbone which will extract feature maps
:param pretrained_backbone: if True, return pretrained backbone of resnet
:return: instance of FasterRCNN
"""
if backbone_n not in names:
raise Exception('Wrong backbone name')
backbone = torchvision.models.detection.backbone_utils.resnet_fpn_backbone(
backbone_n, pretrained=pretrained_backbone)
# set out channels for FasterRCNN
backbone.out_channels = 256
# define custom anchors for RPN
anchor_generator = AnchorGenerator(sizes=cfg.ANCHOR.SCALES,
aspect_ratios=cfg.ANCHOR.RATIOS)
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# define model
model = FasterRCNN(backbone,
num_classes=n_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def fasterrcnn_resnet_fpn(pretrained=False,
progress=True,
resnet='resnet50',
num_classes=91,
pretrained_backbone=True,
**kwargs):
if pretrained:
# no need to download the backbone if pretrained is set
pretrained_backbone = False
backbone = resnet_fpn_backbone(resnet, pretrained_backbone)
model = FasterRCNN(backbone, num_classes, **kwargs)
if pretrained:
target_url = model_urls['fasterrcnn_' + resnet + '_fpn_coco']
state_dict = load_state_dict_from_url(target_url, progress=progress)
model.load_state_dict(state_dict)
return model