| 19 | |
| 20 | def cv2_demo(net, transform): |
| 21 | def predict(frame): |
| 22 | height, width = frame.shape[:2] |
| 23 | x = torch.from_numpy(transform(frame)[0]).permute(2, 0, 1) |
| 24 | x = Variable(x.unsqueeze(0)) |
| 25 | y = net(x) # forward pass |
| 26 | detections = y.data |
| 27 | # scale each detection back up to the image |
| 28 | scale = torch.Tensor([width, height, width, height]) |
| 29 | for i in range(detections.size(1)): |
| 30 | j = 0 |
| 31 | while detections[0, i, j, 0] >= 0.6: |
| 32 | pt = (detections[0, i, j, 1:] * scale).cpu().numpy() |
| 33 | cv2.rectangle(frame, |
| 34 | (int(pt[0]), int(pt[1])), |
| 35 | (int(pt[2]), int(pt[3])), |
| 36 | COLORS[i % 3], 2) |
| 37 | cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), |
| 38 | FONT, 2, (255, 255, 255), 2, cv2.LINE_AA) |
| 39 | j += 1 |
| 40 | return frame |
| 41 | |
| 42 | # start video stream thread, allow buffer to fill |
| 43 | print("[INFO] starting threaded video stream...") |