def __init__(self,
obj_n_classes,
part_n_classes,
fusion_thresh=0.9,
use_attention=True):
super(JointDetector, self).__init__()
self.fusion_thresh = fusion_thresh
self.use_attention = use_attention
print(
'Creating JointDetector(fusion_thresh=%.2f) for %d Object, %d Part classes...'
% (self.fusion_thresh, obj_n_classes, part_n_classes))
self.object_detector = fasterrcnn_resnet50_fpn(pretrained=True)
self.part_detector = fasterrcnn_resnet50_fpn(pretrained=True)
in_features_obj_det = self.object_detector.roi_heads.box_predictor.cls_score.in_features # 1024
in_features_part_det = self.part_detector.roi_heads.box_predictor.cls_score.in_features # 1024
in_features = in_features_obj_det + in_features_part_det
self.object_detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, obj_n_classes)
self.part_detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, part_n_classes)
self.transform = self.object_detector.transform # both detectors have save transforms so anyone can be used
if self.use_attention:
print(
'Using Attention Layer to compute object-part score for feature fusion!'
)
self.attention_layer = nn.Sequential(nn.Linear(in_features, 1))
def _fasterrcnn_resnet_fpn(backbone='resnet50',
num_classes=91,
pretrained_backbone=True,
**kwargs):
if import_error is not None:
raise import_error
from torchvision.models.detection.faster_rcnn import FasterRCNN, resnet_fpn_backbone
backbone = resnet_fpn_backbone(backbone, pretrained_backbone)
model = FasterRCNN(backbone, num_classes, **kwargs)
detection.fasterrcnn_resnet50_fpn()
return model
def build_model(self):
model = fasterrcnn_resnet50_fpn(pretrained=False)
num_classes = 20
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
self.model = model.to(self.device)
def get_FasterRCNN_model(num_classes):
model = fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
def __init__(self, num_classes=det_config.NUM_CLASSES):
super(FasterRCNN, self).__init__()
self.model = fasterrcnn_resnet50_fpn(pretrained=True)
self.num_classes = num_classes
self.in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = FastRCNNPredictor(
self.in_features, self.num_classes)
def predict(img_path, threshold):
model = fasterrcnn_resnet50_fpn(pretrained=True)
model = model.cuda()
model.eval()
try:
img = Image.open(img_path)
except:
print("Image not found......exiting")
exit(0)
trf = transforms.Compose(
[transforms.Resize((1024, 1024)),
transforms.ToTensor()])
img_ = trf(img)
img_ = img_.cuda()
img_ = img_.unsqueeze(0)
with torch.no_grad():
predictions = model(img_)
pred_scores = predictions[0]['scores'].cpu().numpy()
pred_boxes = predictions[0]['boxes'].cpu().numpy()
pred_classes = predictions[0]['labels'].cpu().numpy()
pred_labels = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in pred_classes]
del pred_classes
output = []
for i, val in enumerate(pred_scores):
if val >= threshold:
output.append(((pred_boxes[i][0], pred_boxes[i][1]),
(pred_boxes[i][2],
pred_boxes[i][3]), pred_labels[i], val))
return output
def create_faster_rcnn(
pretrained: bool = False, trainable_backbone_layers: int = 5
) -> FasterRCNN:
"""Creates Faster R-CNN model able to detect objects of
a single class.
Parameters
----------
pretrained : bool
Flag to create a pretrained model on the ImageNet dataset.
trainable_backbone_layers : int
Number of backbone layers to train.
Returns
-------
FasterRCNN
Returns Faster R-CNN model.
"""
model = fasterrcnn_resnet50_fpn(
pretrained=pretrained, trainable_backbone_layers=trainable_backbone_layers
)
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
def track(
video_path: str,
output_path: str = "out.json",
score_threshold: float = 0.5,
class_index: int = 1, # Track people by default
) -> None:
"""Track the objects for a specific class in a given video"""
# Initialization
device = "cuda" if torch.cuda.is_available() else "cpu"
model = fasterrcnn_resnet50_fpn(pretrained=True).eval().to(device)
reader = imageio.get_reader(video_path)
# Tracking loop
active_tracklets = []
finished_tracklets = []
prev_boxes = []
for i, frame in enumerate(reader):
height, width = frame.shape[:2]
# Detection
x = to_tensor(frame).to(device)
result = model(x[None])[0]
# Feature extraction: (x1, y1, x2, y2) in image coordinates
# where class == class_index and score > score_threshold
mask = torch.logical_and(result["labels"] == class_index,
result["scores"] > score_threshold)
boxes = result["boxes"][mask].data.cpu().numpy() / np.array(
[width, height, width, height])
prev_indices = []
boxes_indices = []
if len(boxes) > 0 and len(prev_boxes) > 0:
# Pairwise cost: euclidean distance between boxes
cost = np.linalg.norm(prev_boxes[:, None] - boxes[None], axis=-1)
# Object matching
prev_indices, boxes_indices = linear_sum_assignment(cost)
# Add matches to active tracklets
for prev_idx, box_idx in zip(prev_indices, boxes_indices):
active_tracklets[prev_idx]["boxes"].append(
np.round(boxes[box_idx], 3).tolist())
# Finalize lost tracklets
lost_indices = set(range(len(active_tracklets))) - set(prev_indices)
for lost_idx in sorted(lost_indices, reverse=True):
finished_tracklets.append(active_tracklets.pop(lost_idx))
# Activate new tracklets
new_indices = set(range(len(boxes))) - set(boxes_indices)
for new_idx in new_indices:
active_tracklets.append({
"start":
i,
"boxes": [np.round(boxes[new_idx], 3).tolist()]
})
# "Predict" next frame for comparison
prev_boxes = np.array(
[tracklet["boxes"][-1] for tracklet in active_tracklets])
with open(output_path, "w") as f:
f.write(
json.dumps({
"fps": reader.get_meta_data()["fps"],
"tracklets": finished_tracklets + active_tracklets,
}))
def _fasterrcnn_resnet50(device: str, model_weights: str = None):
"""Creates a FasterRCNN with a ResNet50 backbone."""
model = fasterrcnn_resnet50_fpn()
if model_weights:
checkpoint = torch.load(model_weights, map_location=device)
model.load_state_dict(checkpoint["state_dict"])
return model.eval().to(device)
def __init__(self, num_classes: int = 2, pretrained: bool = True):
super(CoffeAI, self).__init__()
self.num_classes = num_classes + 1 # add background class
self.model = detection.fasterrcnn_resnet50_fpn(pretrained=pretrained)
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = \
detection.faster_rcnn.FastRCNNPredictor(in_features, self.num_classes)
def _build_model(self) -> nn.Module:
model = fasterrcnn_resnet50_fpn(pretrained=True)
# Replace the classifier with a new two-class classifier. There are
# only two "classes": pedestrian and background.
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
def _end_setup(self, use_lists=False):
self.detection_net = fasterrcnn_resnet50_fpn(
pretrained=False,
pretrained_backbone=False,
num_classes=self.options.num_classes).to(self.device)
self.iou_loss = 0
self._create_optimizer([self.detection_net], True, "detection")
super()._end_setup(use_lists)
def get_net():
model = fasterrcnn_resnet50_fpn(pretrained=True,
box_score_thresh=0.8,
min_size=480,
max_size=720)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 3) # 目标数
return model
def get_faster_rcnn():
# load an instance segmentation model pre-trained pre-trained on COCO
model = 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 = faster_rcnn.FastRCNNPredictor(
in_features, 3)
return model
def __init__(self, batch, device):
from torchvision.models.detection import fasterrcnn_resnet50_fpn
import torch
# from torchvision import transforms
super(FasterRCNNTorch, self).__init__()
self.model = fasterrcnn_resnet50_fpn(pretrained=True)
self.model.eval()
self.batch = batch
self.device = torch.device(device)
self.model.to(self.device)
def rcnnInit1():
global model1
model1 = fasterrcnn_resnet50_fpn(
num_classes=num_classes, pretrained_backbone=False, pretrained=False)
try:
model1.load_state_dict(torch.load('/workENV/src/frankNet/RCNN1.pt'))
except Exception:
model1.load_state_dict(torch.load('src/frankNet/RCNN1.pt'))
model1.eval()
print("Load RCNN1 Complete!")
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Load pretrained model
self.model = fasterrcnn_resnet50_fpn(num_classes=91)
if self.pretrain_model is not None:
state_dict = torch.load(self.pretrain_model)
else:
state_dict = load_state_dict_from_url(FasterRCNN.PRETRAIN_URL)
self.model.load_state_dict(state_dict)
self.model.to(self.device)
self.model.eval()
def __init__(self, arch, classnames, pretrained=False):
super().__init__()
assert arch == 'fasterrcnn_resnet50_fpn'
pretrained = pretrained == 'coco'
if classnames is None:
from ..data.coco import COCO_CLASSNAMES
self.classnames = COCO_CLASSNAMES
else:
self.classnames = classnames
self.num_classes = len(self.classnames) + 1
self.model = fasterrcnn_resnet50_fpn(pretrained=pretrained, num_classes=self.num_classes)
def __init__(self):
super(FasterRCNNModel, self).__init__()
self.mask_criterion = nn.BCEWithLogitsLoss()
self.faster_rcnn = fasterrcnn_resnet50_fpn(
num_classes=9, pretrained_backbone=False)
self.deconv1 = nn.ConvTranspose2d(256, 128, 2, 2, 0)
self.norm1 = nn.BatchNorm2d(128)
self.deconv2 = nn.ConvTranspose2d(128, 64, 2, 2, 0)
self.norm2 = nn.BatchNorm2d(64)
self.conv = nn.Conv2d(64, 1, 1, 1, 0)
self.relu = nn.ReLU()
def __init__(
self,
learning_rate: float = 0.0001,
num_classes: int = 91,
pretrained: bool = False,
pretrained_backbone: bool = True,
trainable_backbone_layers: int = 3,
replace_head: bool = True,
**kwargs,
):
"""
PyTorch Lightning implementation of `Faster R-CNN: Towards Real-Time Object Detection with
Region Proposal Networks <https://arxiv.org/abs/1506.01497>`_.
Paper authors: Shaoqing Ren, Kaiming He, Ross Girshick, Jian Sun
Model implemented by:
- `Teddy Koker <https://github.com/teddykoker>`
During training, the model expects both the input tensors, as well as targets (list of dictionary), containing:
- boxes (`FloatTensor[N, 4]`): the ground truth boxes in `[x1, y1, x2, y2]` format.
- labels (`Int64Tensor[N]`): the class label for each ground truh box
CLI command::
# PascalVOC
python faster_rcnn.py --gpus 1 --pretrained True
Args:
learning_rate: the learning rate
num_classes: number of detection classes (including background)
pretrained: if true, returns a model pre-trained on COCO train2017
pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet
trainable_backbone_layers: number of trainable resnet layers starting from final block
"""
super().__init__()
model = fasterrcnn_resnet50_fpn(
# num_classes=num_classes,
pretrained=pretrained,
pretrained_backbone=pretrained_backbone,
trainable_backbone_layers=trainable_backbone_layers,
)
if replace_head:
in_features = model.roi_heads.box_predictor.cls_score.in_features
head = faster_rcnn.FastRCNNPredictor(in_features, num_classes)
model.roi_heads.box_predictor = head
else:
assert num_classes == 91, "replace_head must be true to change num_classes"
self.model = model
self.learning_rate = learning_rate
def __init__(self,
n_channels=3,
n_classes=21,
softmax_out=False,
resnet_type=101,
pretrained=False):
super(FasterRCNN, self).__init__()
self.resnet_type = resnet_type
self.n_channels = n_channels
self.n_classes = n_classes
self.pretrained = pretrained
# Input conv is applied to convert the input to 3 ch depth
self.inconv = None
if n_channels != 3:
self.inconv = FwdConv(n_channels, 3, kernel_size=1, padding=0)
# Pre-trained model needs to be an identical network
if pretrained:
self.body = fasterrcnn_resnet50_fpn(pretrained=pretrained,
num_classes=91,
min_size=512)
# Reset output
if n_classes != 91:
self.body.roi_heads.box_predictor.cls_score = nn.Linear(
in_features=1024, out_features=n_classes, bias=True)
self.body.roi_heads.box_predictor.bbox_pred = nn.Linear(
in_features=1024, out_features=4 * n_classes, bias=True)
else:
self.body = fasterrcnn_resnet50_fpn(pretrained=pretrained,
num_classes=n_classes,
min_size=512)
# Softmax alternative
self.has_softmax = softmax_out
if softmax_out:
self.softmax = nn.Softmax2d()
else:
self.softmax = None
def __init__(self, num_labels, img_sz, pretrained=True):
super().__init__()
self.model = fasterrcnn_resnet50_fpn(pretrained=False,
progress=True,
num_classes=num_labels,
pretrained_backbone=pretrained,
min_size=img_sz,
max_size=img_sz)
self.subloss_names = [
'total_loss', 'loss_box_reg', 'loss_classifier', 'loss_objectness',
'loss_rpn_box_reg'
]
def get_object_detection_model(num_classes, ispretrain=True):
# load an object detection model pre-trained on COCO
model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
# replace the classifier with a new one
num_classes = num_classes # 11 class digit + background
# 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 = detection.faster_rcnn.FastRCNNPredictor(
in_features, num_classes)
return model
def main():
process = psutil.Process(os.getpid())
model = fasterrcnn_resnet50_fpn(pretrained=True)
# model.eval() # Gives no issues
for i in range(10):
size = random.randint(700, 900)
images = torch.rand([1, 3, size, 800])
# images = torch.rand([1, 3, size, size]) # Gives no issues
targets = [{'boxes': torch.tensor([[10., 20., 30., 40.]]), 'labels': torch.tensor([1])}]
model(images, targets)
gc.collect()
print("Current memory: ", process.memory_info()[0] / float(2**20))
def model_dispenser():
model = fasterrcnn_resnet50_fpn(pretrained=True)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
head = FastRCNNPredictor(in_features, num_classes=config.CLASSES)
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = head
return model
def __init__(self,
num_channels=3,
num_classes=91,
pretrained=False,
min_size=512):
super(FasterRCNN, self).__init__()
self.num_channels = num_channels
self.num_classes = num_classes
self.pretrained = pretrained
self.min_size = min_size
# Input conv is applied to convert the input to 3 ch depth
self.inconv = None
if num_channels != 3:
self.inconv = nn.Sequential()
#transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std)
#self.inconv.add_module("transform", GeneralizedRCNNTransform(min_size, min_size, image_mean, image_std))
self.inconv.add_module(
"conv0", nn.Conv2d(num_channels, 3, 1, stride=1, padding=0))
self.inconv.add_module("bn0", nn.BatchNorm2d(3))
self.inconv.add_module("relu", nn.ReLU(inplace=True))
# Pre-trained model needs to be an identical network
if pretrained:
self.body = fasterrcnn_resnet50_fpn(pretrained=pretrained,
num_classes=91,
min_size=min_size)
# Reset output
if num_classes != 91:
in_features = self.body.roi_heads.box_predictor.cls_score.in_features
self.body.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
else:
self.body = fasterrcnn_resnet50_fpn(pretrained=pretrained,
num_classes=num_classes,
min_size=min_size)
def __init__(self, pretrained=True, nrois=10):
super(RPN,self).__init__()
model = fasterrcnn_resnet50_fpn(pretrained=True).eval()
self.transform = STRGTransform(model.transform.min_size,
model.transform.max_size,
0,0) #copy.deepcopy(model.transform)
self.backbone = copy.deepcopy(model.backbone)
self.rpn = copy.deepcopy(model.rpn)
# self.eaget_outputs = copy.deepcopy(model.eaget_outputs)
self.roi_heads = copy.deepcopy(model.roi_heads)
self.rpn._pre_nms_top_n = {'training':3*nrois, 'testing':3*nrois}
self.rpn._post_nms_top_n = {'training':nrois, 'testing':nrois}
self.rpn.fg_bg_sampler.positive_fraction = 1.0
del model
def get_model(num_classes, backbone=None):
'''
Model function to output network according to arguments.
Args:
num_classes: number of classes(total_classes+1 for background)
backbone: to design network with other backbone, default backbone
of faster RCNN is resnet50.
Returns:
model.
'''
if backbone == 'mobile_net':
net = mobilenet_v2(pretrained=True)
backbone_ft = net.features
backbone_ft.out_channels = 1280
elif backbone == 'vgg19':
net = vgg19(pretrained=True)
backbone_ft = net.features
backbone_ft.out_channels = 512
# https://stackoverflow.com/questions/58362892/resnet-18-as-backbone-in-faster-r-cnn
elif backbone == 'resnet101':
net = resnet101(pretrained=True)
modules = list(net.children())[:-1]
backbone_ft = nn.Sequential(*modules)
backbone_ft.out_channels = 2048
if backbone is None:
model = fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
# print(in_features) = 1024
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
return model
else:
anchor_gen = AnchorGenerator(sizes=((32, 64, 128), ))
# featmap_names = [0] gives list index out of range error.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone_ft,
num_classes,
rpn_anchor_generator=anchor_gen,
box_roi_pool=roi_pooler)
return model
def test_fasterrcnn_resnet50_fpn_frozen_layers(self):
# we know how many initial layers and parameters of the network should
# be frozen for each trainable_backbone_layers paramter value
# i.e all 53 params are frozen if trainable_backbone_layers=0
# ad first 24 params are frozen if trainable_backbone_layers=2
expected_frozen_params = {0: 53, 1: 43, 2: 24, 3: 11, 4: 1, 5: 0}
for train_layers, exp_froz_params in expected_frozen_params.items():
model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False,
num_classes=91, pretrained_backbone=False,
trainable_backbone_layers=train_layers)
# boolean list that is true if the param at that index is frozen
is_frozen = [not parameter.requires_grad for _, parameter in model.named_parameters()]
# check that expected initial number of layers are frozen
self.assertTrue(all(is_frozen[:exp_froz_params]))
def __init__(self, model_path):
num_classes = 2
self.to_tensor = ToTensor()
self.device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
self.model = fasterrcnn_resnet50_fpn()
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
self.model.to(self.device)
self.model.load_state_dict(
torch.load(model_path, map_location=self.device))
self.model.eval()
