-
Notifications
You must be signed in to change notification settings - Fork 0
/
pred_yolo.py
66 lines (52 loc) · 1.87 KB
/
pred_yolo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import numpy as np
import cv2
import os
from pathlib import Path
import pandas as pd
def get_preds(net, img_path):
# load our input image and grab its spatial dimensions
image = cv2.imread(img_path)
(H, W) = image.shape[:2]
# determine only the *output* layer names that we need from YOLO
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# construct a blob from the input image and then perform a forward
# pass of the YOLO object detector, giving us our bounding boxes and
# associated probabilities
blob = cv2.dnn.blobFromImage(image,
1 / 255.0,
(416, 416),
swapRB=True,
crop=False)
net.setInput(blob)
layerOutputs = net.forward(ln)
scores = None
# loop over each of the layer outputs
for output in layerOutputs:
# loop over each of the detections
for detection in output:
# extract the class ID and confidence (i.e., probability) of
# the current object detection
scores_new = detection[5:]
scores = scores_new if scores is None else scores + scores_new
return scores
# load the COCO class labels our YOLO model was trained on
labelsPath = os.path.abspath('./openimages.names')
LABELS = open(labelsPath).read().strip().split("\n")
# derive the paths to the YOLO weights and model configuration
weightsPath = os.path.abspath('./yolov3-openimages.weights')
configPath = os.path.abspath('./yolov3-openimages.cfg')
# load our YOLO object detector
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
p = Path(r'D:\Interactive Video Retrieval\keyframes\keyframes')
res_dict = {}
n_keyframes = 108645
counter = 0
for img_path in p.glob('**/*.png'):
counter += 1
print(counter, '/', n_keyframes)
res = get_preds(net=net, img_path=str(img_path.resolve()))
res_dict[img_path.stem] = res.reshape(-1, )
df = pd.DataFrame.from_dict(res_dict, orient="index")
df.to_csv(f'./yolo_preds.csv')
print(df)