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)