class TorchDetector:
"""
Torch object detector
"""
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 stop(self):
"""
Destruction
"""
def detectObjects(self, img) -> List[e.DetectedObject]:
"""
Implementation of detector interface
"""
wsize = 1600
hsize = 800
_pretransform = A.Compose([
A.Resize(hsize, wsize),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
image_tensor = _pretransform(image=img)['image']
tstart = time.time()
outputs = self._model.forward(
image_tensor.unsqueeze(0).float().to(device=self._device))
classes = outputs[0]['labels'].detach().cpu().numpy()
scores = outputs[0]['scores'].detach().cpu().numpy()
boxes = outputs[0]['boxes'].detach().cpu().numpy()
self._logger.debug(
f'Torch model inferring time: {time.time() - tstart}')
result = zip(classes, scores, boxes)
h, w, _ = img.shape
wscale = w / wsize
hscale = h / hsize
#print(f'h,w:{h},{w}; wsc,hsc:{wscale},{hscale}')
#print(list(result))
return ObjectDetector.getDetectedObjectsCollection(
result, hscale, wscale, self._threshold, False)
class TrashLiteModel(nn.Module):
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 forward(self, images, targets=None):
return self.model.forward(images, targets)