main.py 4.05 KB
Newer Older
Pavlo Beylin's avatar
Pavlo Beylin committed
1
from PIL import Image
Pavlo Beylin's avatar
Pavlo Beylin committed
2
3
4
import torch
import cv2
import time
Pavlo Beylin's avatar
Pavlo Beylin committed
5
6
import matplotlib
matplotlib.use('TkAgg')
Pavlo Beylin's avatar
Pavlo Beylin committed
7
8
9
import matplotlib.pyplot as plt

# Model
Pavlo Beylin's avatar
Pavlo Beylin committed
10
11
12
13
from torchvision.transforms import transforms

from patch_transformer import PatchTransformer, PatchApplier

Pavlo Beylin's avatar
Pavlo Beylin committed
14
model = torch.hub.load('ultralytics/yolov5', 'yolov5l')  # or yolov5m, yolov5l, yolov5x, cu
Pavlo Beylin's avatar
Pavlo Beylin committed
15
16


Pavlo Beylin's avatar
Pavlo Beylin committed
17
18
19
20
# model = torch.hub.load('ultralytics/yolov3', 'yolov3')


def show(img):
Pavlo Beylin's avatar
Pavlo Beylin committed
21
    plt.imshow(img.detach().cpu())
Pavlo Beylin's avatar
Pavlo Beylin committed
22
23
    plt.show()

Pavlo Beylin's avatar
Pavlo Beylin committed
24

Pavlo Beylin's avatar
Pavlo Beylin committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
classes = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus",
           "train", "truck", "boat", "traffic light", "fire hydrant",
           "stop sign", "parking meter", "bench", "bird", "cat", "dog",
           "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
           "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
           "skis", "snowboard", "sports ball", "kite", "baseball bat",
           "baseball glove", "skateboard", "surfboard", "tennis racket",
           "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
           "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
           "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant",
           "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote",
           "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
           "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"]

Pavlo Beylin's avatar
Pavlo Beylin committed
39
40
41
PATH = "cat_patch0.jpg"
PATCH_SIZE = 100

Pavlo Beylin's avatar
Pavlo Beylin committed
42
43
44
45
46
47
48
49
50
51
52

def debug_preds():
    detected_classes = [int(results.pred[0][i][-1]) for i in range(0, len(results.pred[0]))]
    print(detected_classes)
    for det in results.pred[0]:
        if int(det[-1]) == 0:  # person
            print("Person ({}):".format(float(det[-2])))
            print("x1:y1 : {}:{}".format(float(det[0]), float(det[1])))
            print("x2:y2 : {}:{}".format(float(det[2]), float(det[3])))


Pavlo Beylin's avatar
Pavlo Beylin committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
# from https://github.com/wangzh0ng/adversarial_yolo2
def read_image(path):
    """
    Read an input image to be used as a patch

    :param path: Path to the image to be read.
    :return: Returns the transformed patch as a pytorch Tensor.
    """
    patch_img = Image.open(path).convert('RGB')
    tf = transforms.Resize((PATCH_SIZE, PATCH_SIZE))
    patch_img = tf(patch_img)
    tf = transforms.ToTensor()

    return tf(patch_img)




if __name__ == "__main__":
    # init
    patch_transformer = PatchTransformer().cuda()
    patch_applier = PatchApplier().cuda()

Pavlo Beylin's avatar
Pavlo Beylin committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

    # set start time to current time
    start_time = time.time()

    # displays the frame rate every 2 second
    display_time = 2

    # Set primary FPS to 0
    fps = 0

    # we create the video capture object cap
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise IOError("We cannot open webcam")

Pavlo Beylin's avatar
Pavlo Beylin committed
91
92
93
94
95
96
    patch = read_image(PATH)
    patch_np = torch.transpose(patch.T, 0, 1).numpy()

    img_size_x = 640
    img_size_y = 480

Pavlo Beylin's avatar
Pavlo Beylin committed
97
98
99
100
    while True:
        ret, frame = cap.read()
        # resize our captured frame if we need
        frame = cv2.resize(frame, None, fx=1.0, fy=1.0, interpolation=cv2.INTER_AREA)
Pavlo Beylin's avatar
Pavlo Beylin committed
101
102
103
104
105
106
107
108
109
110
        # cv2.imshow("Web cam input", frame)

        # transform patch
        trans_patch = patch_transformer(patch.cuda(), torch.ones([1, 14, 5]).cuda(), img_size_x, img_size_y,
                                        do_rotate=True, rand_loc=True)
        trans_patch_np = torch.transpose(trans_patch[0][0].T, 0, 1).detach().cpu().numpy()
        # cv2.imshow("patch", trans_patch_np)

        # apply patch
        frame = patch_applier(frame, trans_patch_np)
Pavlo Beylin's avatar
Pavlo Beylin committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

        # detect object on our frame
        results = model(frame.copy())
        # debug_preds()

        # show us frame with detection
        cv2.imshow("img", results.render()[0])
        if cv2.waitKey(25) & 0xFF == ord("q"):
            cv2.destroyAllWindows()
            break

        # calculate FPS
        fps += 1
        TIME = time.time() - start_time
        if TIME > display_time:
            print("FPS:", fps / TIME)
            fps = 0
            start_time = time.time()

    cap.release()
    cv2.destroyAllWindows()