what.models.detection.ssd.ssd.predictor

 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:
 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)
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):
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]