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]