what.models.detection.yolox

This module implements YOLOX object detection model.

what.models.detection.yolox.yolox_x

what.models.detection.yolox.yolox_m

what.models.detection.yolox.yolox_l

what.models.detection.yolox.yolox_s

View Source

 1r'''
 2This module implements YOLOX object detection model.
 3
 4<br />
 5
 6## what.models.detection.yolox.yolox_x
 7## what.models.detection.yolox.yolox_m
 8## what.models.detection.yolox.yolox_l
 9## what.models.detection.yolox.yolox_s
10
11'''
12
13from what.models.detection.yolox.yolox_x import YOLOX_X
14from what.models.detection.yolox.yolox_m import YOLOX_M
15from what.models.detection.yolox.yolox_l import YOLOX_L
16from what.models.detection.yolox.yolox_s import YOLOX_S
17
18__all__ = ["YOLOX_X", "YOLOX_M", "YOLOX_L", "YOLOX_S"]

class YOLOX_X: View Source

12class YOLOX_X:
13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-x")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names
28
29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

YOLOX_X(class_names, model_path) View Source

13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-x")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names

def predict(self, image): View Source

29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

class YOLOX_M: View Source

12class YOLOX_M:
13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-m")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names
28
29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

YOLOX_M(class_names, model_path) View Source

13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-m")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names

def predict(self, image): View Source

29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

class YOLOX_L: View Source

12class YOLOX_L:
13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-l")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names
28
29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

YOLOX_L(class_names, model_path) View Source

13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-l")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names

def predict(self, image): View Source

29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

class YOLOX_S: View Source

12class YOLOX_S:
13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-s")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names
28
29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs

YOLOX_S(class_names, model_path) View Source

13    def __init__(self, class_names, model_path):
14        exp = get_exp(None, "yolox-s")
15        self.model = exp.get_model()
16
17        # load the model state dict
18        device = 'gpu' if torch.cuda.is_available() else 'cpu'
19        ckpt = torch.load(model_path, map_location="cpu")
20        self.model.load_state_dict(ckpt["model"])
21
22        self.predictor = Predictor(
23                self.model, exp, COCO_CLASSES, None, None,
24                device, False, False,
25            )
26
27        self.class_names = class_names

def predict(self, image): View Source

29    def predict(self, image):
30        outputs, _ = self.predictor.inference(image)
31        height, width, _ = image.shape
32        scale = min(640 / float(height), 640 / float(width))
33
34        # Run inference
35        class_ids = []
36        boxes = []
37        probs = []
38
39        if outputs[0] is not None:
40            boxes  = outputs[0][:, 0:4].cpu().numpy()
41            class_ids = outputs[0][:, -1].cpu().numpy().astype(np.uint32)
42            probs = (outputs[0][:, 4] * outputs[0][:, 5]).cpu().numpy()
43
44            boxes = np.array([box for box, prob in zip(boxes, probs) if prob >= 0.5 ])
45            probs = probs[probs >= 0.5]
46
47        # From (x1, y1, x2, y2) --> (x, y, w, h)
48        for i, box in enumerate(boxes):
49            box /= scale
50            box[0] /= width
51            box[1] /= height
52            box[2] /= width 
53            box[3] /= height
54
55            box[2] -= box[0]
56            box[3] -= box[1]
57            box[0] += (box[2] / 2)
58            box[1] += (box[3] / 2)
59
60        return image, boxes, class_ids, probs