-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlabel2.py
162 lines (138 loc) · 5.26 KB
/
label2.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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import cv2
import numpy as np
import matplotlib.pyplot as plt
import json
import time
from urllib.error import HTTPError
def save_frame_as_image(frame, frame_id, path_format="mask_{}.png"):
"""Save the given frame as an image file."""
cv2.imwrite(path_format.format(frame_id), frame)
def process_frame(frame):
height, width, channels = frame.shape
blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
net.setInput(blob)
outs = net.forward(output_layers)
class_ids = []
confidences = []
boxes = []
for out in outs:
for detection in out:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.5:
center_x = int(detection[0] * width)
center_y = int(detection[1] * height)
w = int(detection[2] * width)
h = int(detection[3] * height)
x = int(center_x - w / 2)
y = int(center_y - h / 2)
boxes.append([x, y, w, h])
confidences.append(float(confidence))
class_ids.append(class_id)
indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
detections = []
if len(indexes) > 0:
for i in indexes.flatten():
x, y, w, h = boxes[i]
label = str(classes[class_ids[i]])
detections.append({
"label": label,
"confidence": confidences[i],
"box": [x, y, w, h]
})
color = (0, 255, 0)
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
cv2.putText(frame, label, (x, y + 30), font, 1, color, 2)
return detections, frame
labelsPath = "coco.names"
# Load YOLO model
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
# Get class labels
with open(labelsPath, "r") as f:
classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
# Read video
cap = cv2.VideoCapture("path-to-webvid10M/014701_014750/25274.mp4")
# Display image using matplotlib
def show_frame_with_matplotlib(frame):
plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()
results = []
frame_id = 0
frame_rate = cap.get(cv2.CAP_PROP_FPS)
frame_3_sec = int(frame_rate * 3) # Frame at the 3rd second
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Indices of frames to process (first frame, last frame, three middle frames)
frame_indices = [
frame_3_sec, # First frame
total_frames - 1 # Last frame
]
# Three middle frames
interval = (total_frames - 2) // 4
frame_indices.extend([interval, interval * 2, interval * 3])
# Process frames
for frame_id in frame_indices:
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
ret, frame = cap.read()
if not ret:
break
try:
detections, processed_frame = process_frame(frame)
except HTTPError as e:
if e.code == 429:
print("HTTP 429 Error: Too Many Requests. Retrying after delay...")
time.sleep(5) # Wait for 5 seconds before retrying
continue
elif e.code == 500:
print("HTTP 500 Error: Internal Server Error. Saving results and stopping...")
# Save current results to JSON file
with open("mask.json", "w") as f:
json.dump(results, f, indent=4)
break
else:
raise
except Exception as e:
if '429' in str(e):
print("HTTP 429 Error: Too Many Requests. Retrying after delay...")
time.sleep(5) # Wait for 5 seconds before retrying
continue
elif '500' in str(e):
print("HTTP 500 Error: Internal Server Error. Saving results and stopping...")
# Save current results to JSON file
with open("mask.json", "w") as f:
json.dump(results, f, indent=4)
break
else:
raise
results.append({
"frame_id": frame_id,
"detections": detections
})
plt.imshow(cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB))
plt.pause(0.1)
cap.release()
plt.close()
# Save results to JSON file
with open("mask.json", "w") as f:
json.dump(results, f, indent=4)
# Setup for capturing and saving frames
cap = cv2.VideoCapture("path-to-webvid10M/014701_014750/25274.mp4")
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Calculate indices for first, last, and three middle frames
frame_indices = [0, total_frames - 1] # first and last frames
middle_frame_step = total_frames // 4
middle_indices = [middle_frame_step, 2 * middle_frame_step, 3 * middle_frame_step]
frame_indices.extend(middle_indices)
# Process and save specified frames
for i, frame_index in enumerate(frame_indices):
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
ret, frame = cap.read()
if ret:
save_frame_as_image(frame, i + 1) # Save frame as mask_{i+1}.png
else:
print(f"Failed to retrieve frame at index {frame_index}")
cap.release()