what.models.detection.ssd.ssd.predictor API documentation

# what.models.detection.ssd.ssd.predictor

View Source

     1import torch
     2
     3from ..utils import box_utils
     4from .preprocessing import PredictionTransform
     5
     6from what.models.detection.utils.time_utils import Timer
     7
     8class Predictor:
     9    def __init__(self, net, size, mean=0.0, std=1.0, nms_method=None,
    10                 iou_threshold=0.45, filter_threshold=0.01, candidate_size=200, sigma=0.5, device=None):
    11        self.net = net
    12        self.transform = PredictionTransform(size, mean, std)
    13        self.iou_threshold = iou_threshold
    14        self.filter_threshold = filter_threshold
    15        self.candidate_size = candidate_size
    16        self.nms_method = nms_method
    17
    18        self.sigma = sigma
    19        if device:
    20            self.device = device
    21        else:
    22            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    23
    24        self.net.to(self.device)
    25        self.net.eval()
    26
    27        self.timer = Timer()
    28
    29    def predict(self, image, top_k=-1, prob_threshold=None):
    30        cpu_device = torch.device("cpu")
    31        height, width, _ = image.shape
    32        image = self.transform(image)
    33        images = image.unsqueeze(0)
    34        images = images.to(self.device)
    35        images.requires_grad = True
    36        # with torch.no_grad():
    37        self.timer.start()
    38        scores, boxes = self.net.forward(images)
    39        print("FPS: ", int(1.0 / self.timer.end()))
    40        boxes = boxes[0]
    41        scores = scores[0]
    42        if not prob_threshold:
    43            prob_threshold = self.filter_threshold
    44        # this version of nms is slower on GPU, so we move data to CPU.
    45        boxes = boxes.to(cpu_device)
    46        scores = scores.to(cpu_device)
    47        picked_box_probs = []
    48        picked_labels = []
    49        for class_index in range(1, scores.size(1)):
    50            probs = scores[:, class_index]
    51            mask = probs > prob_threshold
    52            probs = probs[mask]
    53            if probs.size(0) == 0:
    54                continue
    55            subset_boxes = boxes[mask, :]
    56            box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
    57            box_probs = box_utils.nms(box_probs, self.nms_method,
    58                                      score_threshold=prob_threshold,
    59                                      iou_threshold=self.iou_threshold,
    60                                      sigma=self.sigma,
    61                                      top_k=top_k,
    62                                      candidate_size=self.candidate_size)
    63            picked_box_probs.append(box_probs)
    64            picked_labels.extend([class_index] * box_probs.size(0))
    65        if not picked_box_probs:
    66            return torch.tensor([]), torch.tensor([]), torch.tensor([]), torch.tensor([])
    67        picked_box_probs = torch.cat(picked_box_probs)
    68
    69        # (x1, y1, x2, y2) --> (c1, c2, w, h) (0.0, 1.0)
    70        box_w  = picked_box_probs[:, 2] - picked_box_probs[:, 0]
    71        box_h = picked_box_probs[:, 3] - picked_box_probs[:, 1]
    72        picked_box_probs[:, 0] += box_w / 2
    73        picked_box_probs[:, 1] += box_h / 2
    74        picked_box_probs[:, 2] = box_w
    75        picked_box_probs[:, 3] = box_h
    76
    77        return images, picked_box_probs[:, :4], torch.tensor(picked_labels), picked_box_probs[:, 4]

class Predictor: View Source

     9class Predictor:
    10    def __init__(self, net, size, mean=0.0, std=1.0, nms_method=None,
    11                 iou_threshold=0.45, filter_threshold=0.01, candidate_size=200, sigma=0.5, device=None):
    12        self.net = net
    13        self.transform = PredictionTransform(size, mean, std)
    14        self.iou_threshold = iou_threshold
    15        self.filter_threshold = filter_threshold
    16        self.candidate_size = candidate_size
    17        self.nms_method = nms_method
    18
    19        self.sigma = sigma
    20        if device:
    21            self.device = device
    22        else:
    23            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    24
    25        self.net.to(self.device)
    26        self.net.eval()
    27
    28        self.timer = Timer()
    29
    30    def predict(self, image, top_k=-1, prob_threshold=None):
    31        cpu_device = torch.device("cpu")
    32        height, width, _ = image.shape
    33        image = self.transform(image)
    34        images = image.unsqueeze(0)
    35        images = images.to(self.device)
    36        images.requires_grad = True
    37        # with torch.no_grad():
    38        self.timer.start()
    39        scores, boxes = self.net.forward(images)
    40        print("FPS: ", int(1.0 / self.timer.end()))
    41        boxes = boxes[0]
    42        scores = scores[0]
    43        if not prob_threshold:
    44            prob_threshold = self.filter_threshold
    45        # this version of nms is slower on GPU, so we move data to CPU.
    46        boxes = boxes.to(cpu_device)
    47        scores = scores.to(cpu_device)
    48        picked_box_probs = []
    49        picked_labels = []
    50        for class_index in range(1, scores.size(1)):
    51            probs = scores[:, class_index]
    52            mask = probs > prob_threshold
    53            probs = probs[mask]
    54            if probs.size(0) == 0:
    55                continue
    56            subset_boxes = boxes[mask, :]
    57            box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
    58            box_probs = box_utils.nms(box_probs, self.nms_method,
    59                                      score_threshold=prob_threshold,
    60                                      iou_threshold=self.iou_threshold,
    61                                      sigma=self.sigma,
    62                                      top_k=top_k,
    63                                      candidate_size=self.candidate_size)
    64            picked_box_probs.append(box_probs)
    65            picked_labels.extend([class_index] * box_probs.size(0))
    66        if not picked_box_probs:
    67            return torch.tensor([]), torch.tensor([]), torch.tensor([]), torch.tensor([])
    68        picked_box_probs = torch.cat(picked_box_probs)
    69
    70        # (x1, y1, x2, y2) --> (c1, c2, w, h) (0.0, 1.0)
    71        box_w  = picked_box_probs[:, 2] - picked_box_probs[:, 0]
    72        box_h = picked_box_probs[:, 3] - picked_box_probs[:, 1]
    73        picked_box_probs[:, 0] += box_w / 2
    74        picked_box_probs[:, 1] += box_h / 2
    75        picked_box_probs[:, 2] = box_w
    76        picked_box_probs[:, 3] = box_h
    77
    78        return images, picked_box_probs[:, :4], torch.tensor(picked_labels), picked_box_probs[:, 4]

Predictor( net, size, mean=0.0, std=1.0, nms_method=None, iou_threshold=0.45, filter_threshold=0.01, candidate_size=200, sigma=0.5, device=None) View Source

    10    def __init__(self, net, size, mean=0.0, std=1.0, nms_method=None,
    11                 iou_threshold=0.45, filter_threshold=0.01, candidate_size=200, sigma=0.5, device=None):
    12        self.net = net
    13        self.transform = PredictionTransform(size, mean, std)
    14        self.iou_threshold = iou_threshold
    15        self.filter_threshold = filter_threshold
    16        self.candidate_size = candidate_size
    17        self.nms_method = nms_method
    18
    19        self.sigma = sigma
    20        if device:
    21            self.device = device
    22        else:
    23            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    24
    25        self.net.to(self.device)
    26        self.net.eval()
    27
    28        self.timer = Timer()

def predict(self, image, top_k=-1, prob_threshold=None): View Source

    30    def predict(self, image, top_k=-1, prob_threshold=None):
    31        cpu_device = torch.device("cpu")
    32        height, width, _ = image.shape
    33        image = self.transform(image)
    34        images = image.unsqueeze(0)
    35        images = images.to(self.device)
    36        images.requires_grad = True
    37        # with torch.no_grad():
    38        self.timer.start()
    39        scores, boxes = self.net.forward(images)
    40        print("FPS: ", int(1.0 / self.timer.end()))
    41        boxes = boxes[0]
    42        scores = scores[0]
    43        if not prob_threshold:
    44            prob_threshold = self.filter_threshold
    45        # this version of nms is slower on GPU, so we move data to CPU.
    46        boxes = boxes.to(cpu_device)
    47        scores = scores.to(cpu_device)
    48        picked_box_probs = []
    49        picked_labels = []
    50        for class_index in range(1, scores.size(1)):
    51            probs = scores[:, class_index]
    52            mask = probs > prob_threshold
    53            probs = probs[mask]
    54            if probs.size(0) == 0:
    55                continue
    56            subset_boxes = boxes[mask, :]
    57            box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
    58            box_probs = box_utils.nms(box_probs, self.nms_method,
    59                                      score_threshold=prob_threshold,
    60                                      iou_threshold=self.iou_threshold,
    61                                      sigma=self.sigma,
    62                                      top_k=top_k,
    63                                      candidate_size=self.candidate_size)
    64            picked_box_probs.append(box_probs)
    65            picked_labels.extend([class_index] * box_probs.size(0))
    66        if not picked_box_probs:
    67            return torch.tensor([]), torch.tensor([]), torch.tensor([]), torch.tensor([])
    68        picked_box_probs = torch.cat(picked_box_probs)
    69
    70        # (x1, y1, x2, y2) --> (c1, c2, w, h) (0.0, 1.0)
    71        box_w  = picked_box_probs[:, 2] - picked_box_probs[:, 0]
    72        box_h = picked_box_probs[:, 3] - picked_box_probs[:, 1]
    73        picked_box_probs[:, 0] += box_w / 2
    74        picked_box_probs[:, 1] += box_h / 2
    75        picked_box_probs[:, 2] = box_w
    76        picked_box_probs[:, 3] = box_h
    77
    78        return images, picked_box_probs[:, :4], torch.tensor(picked_labels), picked_box_probs[:, 4]