demo-ml-pennfudanped

1
import random
2
from torchvision.transforms import functional as F
3

4

5
def _flip_coco_person_keypoints(kps, width):
6
    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
7
    flipped_data = kps[:, flip_inds]
8
    flipped_data[..., 0] = width - flipped_data[..., 0]
9
    # Maintain COCO convention that if visibility == 0, then x, y = 0
10
    inds = flipped_data[..., 2] == 0
11
    flipped_data[inds] = 0
12
    return flipped_data
13

14

15
class Compose(object):
16
    def __init__(self, transforms):
17
        self.transforms = transforms
18

19
    def __call__(self, image, target):
20
        for t in self.transforms:
21
            image, target = t(image, target)
22
        return image, target
23

24

25
class RandomHorizontalFlip(object):
26
    def __init__(self, prob):
27
        self.prob = prob
28

29
    def __call__(self, image, target):
30
        if random.random() < self.prob:
31
            height, width = image.shape[-2:]
32
            image = image.flip(-1)
33
            bbox = target["boxes"]
34
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
35
            target["boxes"] = bbox
36
            if "masks" in target:
37
                target["masks"] = target["masks"].flip(-1)
38
            if "keypoints" in target:
39
                keypoints = target["keypoints"]
40
                keypoints = _flip_coco_person_keypoints(keypoints, width)
41
                target["keypoints"] = keypoints
42
        return image, target
43

44

45
class ToTensor(object):
46
    def __call__(self, image, target):
47
        image = F.to_tensor(image)
48
        return image, target
49