-
Notifications
You must be signed in to change notification settings - Fork 0
/
insrtructionyfivecropprcnmscr.py
175 lines (139 loc) · 5.75 KB
/
insrtructionyfivecropprcnmscr.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
163
164
165
166
167
168
169
170
171
172
173
174
import torch
import numpy as np
import cv2
import csv
import time
import os
from datetime import datetime
# Load the YOLOv5 model
weights = "yolov5s.pt"
model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights)
# Load the dataset configuration
data = "data/coco128.yaml"
model.yaml = data
# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device).eval()
# OpenCV setup for video capture
cap = cv2.VideoCapture(0) # Use webcam (change the index if you have multiple cameras)
fps = cap.get(cv2.CAP_PROP_FPS)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Define the codec and create a VideoWriter object
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('output.avi', fourcc, fps, (frame_width, frame_height))
# CSV setup for saving detection results
csv_file = 'detection_results.csv'
csv_fields = ['timestamp', 'class', 'confidence', 'x', 'y', 'width', 'height']
csv_output = open(csv_file, 'w')
csv_writer = csv.DictWriter(csv_output, fieldnames=csv_fields)
csv_writer.writeheader()
# Folder setup for saving cropped detections
output_folder = 'detection_crops'
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# Load class labels
class_labels = model.names
# Load classes and their numbers from CSV
class_mapping = {}
with open('class_mapping.csv', 'r') as csv_file:
csv_reader = csv.reader(csv_file)
next(csv_reader) # Skip the header row
for row in csv_reader:
class_mapping[row[0]] = int(row[1])
# Load thresholds from threshold.csv
thresholds = {}
with open('threshold.csv', 'r') as f:
reader = csv.reader(f)
for row in reader:
class_name = row[0]
threshold = int(row[1])
thresholds[class_name] = threshold
# Load starting times for classes from classes.csv
class_start_times = {}
with open('classes.csv', 'r') as f:
reader = csv.reader(f)
for row in reader:
class_name = row[0]
start_time = datetime.strptime(row[1], "%Y-%m-%d %H:%M:%S")
class_start_times[class_name] = start_time
# Initialize variables for counting
counters = {class_name: 0 for class_name in class_mapping.keys()}
# Time variables
start_time = time.time()
previous_time = start_time
# Non-maximum suppression parameters
confidence_threshold = 0.5
nms_threshold = 0.3
# Object detection loop
while True:
ret, frame = cap.read()
if not ret:
break
current_time = time.time()
elapsed_time = current_time - previous_time
if elapsed_time < 0.1:
continue
# Perform object detection
results = model(frame)
# Get detection information
detections = results.pandas().xyxy[0]
# Filter detections for specified classes
filtered_detections = detections[detections['name'].isin(class_mapping.keys())]
# Check if any specified classes are detected
if len(filtered_detections) > 0:
for idx, detection in filtered_detections.iterrows():
class_name = detection['name']
class_number = class_mapping[class_name]
# Check if the class number is 999
if class_number == 999:
# Save detection results in CSV and crop the image
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
csv_writer.writerow({
'timestamp': timestamp,
'class': class_name,
'confidence': detection['confidence'],
'x': detection['xmin'],
'y': detection['ymin'],
'width': detection['xmax'] - detection['xmin'],
'height': detection['ymax'] - detection['ymin']
})
crop = frame[int(detection['ymin']):int(detection['ymax']),
int(detection['xmin']):int(detection['xmax'])]
# Check if the counter exceeds the threshold and the elapsed time exceeds the class start time
if counters[class_name] >= thresholds[class_name] and elapsed_time >= class_start_times[class_name]:
# Save the cropped image with timestamp and location
location = "bus_station_village_district" # Replace with the actual location
image_name = f"{class_name}_{timestamp}_{location}.jpg"
cv2.imwrite(os.path.join(output_folder, image_name), crop)
# Reset the counter
counters[class_name] = 0
# Increment the counter
counters[class_name] += 1
# Display the frame with bounding boxes and labels after applying non-maximum suppression
if len(results.pred) > 0:
pred = results.pred[0]
boxes = pred[:, :4] # Extract the bounding boxes
scores = pred[:, 4] # Extract the confidence scores
labels = pred[:, 5] # Extract the labels
# Apply non-maximum suppression
indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), confidence_threshold, nms_threshold)
for i in indices:
i = i[0]
bbox = boxes[i].int().cpu().numpy() # Convert bbox to int values
label = int(labels[i])
confidence = float(scores[i])
# Draw bounding box and label on the frame
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
cv2.putText(frame, f'{class_labels[label]} {confidence:.2f}', (bbox[0], bbox[1] - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
cv2.imshow('Object Detection', frame)
out.write(frame)
if cv2.waitKey(1) == 27: # Press Esc to exit
break
previous_time = current_time
# Release resources
cap.release()
out.release()
cv2.destroyAllWindows()
csv_output.close()