TKK_E32211602_2/yolo-v3.py

209 lines
7.7 KiB
Python

import os
import cv2
import numpy as np
import requests
import urllib.request
import time
# URL untuk file weights dan konfigurasi YOLOv3-tiny
weights_url_tiny = "https://pjreddie.com/media/files/yolov3-tiny.weights"
config_url_tiny = "https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3-tiny.cfg"
# Path tempat Anda ingin menyimpan file weights dan konfigurasi
weights_path_tiny = "yolo_models/yolov3-tiny.weights"
config_path_tiny = "yolo_models/yolov3-tiny.cfg"
# Buat direktori jika belum ada
os.makedirs("yolo_models", exist_ok=True)
# Fungsi untuk mengunduh file dari URL dan menyimpannya
def download_file(url, save_path):
print(f"Downloading {url}...")
urllib.request.urlretrieve(url, save_path)
print(f"File saved at {save_path}")
# Unduh file weights jika belum ada
if not os.path.exists(weights_path_tiny):
download_file(weights_url_tiny, weights_path_tiny)
# Unduh file konfigurasi jika belum ada
if not os.path.exists(config_path_tiny):
download_file(config_url_tiny, config_path_tiny)
# Load YOLOv3-tiny model and config file
net = cv2.dnn.readNetFromDarknet(config_path_tiny, weights_path_tiny)
# Jika Anda memiliki GPU yang didukung, aktifkan CUDA
if cv2.cuda.getCudaEnabledDeviceCount() > 0:
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
else:
print("CUDA is not available. Running on CPU.")
# Open your laptop's camera
cap = cv2.VideoCapture(1) # Use '0' for built-in webcam, '1' for external webcam
assert cap.isOpened(), "Error accessing webcam"
# Kurangi frame rate kamera
cap.set(cv2.CAP_PROP_FPS, 15) # Set FPS to 15 for reducing load
# Initialize variables
people_in = 0
people_out = 0
next_person_id = 0
tracked_objects = {}
last_sent_people_in = 0
last_sent_people_out = 0
# Function to check if a person crosses the vertical line
def is_crossing_line(prev_centroid, current_centroid, line_position):
return (prev_centroid[0] <= line_position < current_centroid[0]) or (prev_centroid[0] >= line_position > current_centroid[0])
# Function to send data to ESP32
def send_to_esp32(in_count, out_count, total_count):
url = 'http://192.168.249.72/update_counts' # Replace with the IP address of your ESP32
payload = {
'in': in_count,
'out': out_count,
'total': total_count
}
headers = {'Content-Type': 'application/json'}
try:
response = requests.post(url, json=payload, headers=headers)
print(f'Status Code: {response.status_code}, Response: {response.text}')
except requests.exceptions.RequestException as e:
print(f'Error sending data to ESP32: {e}')
while True:
start_time = time.time() # Track the start time
ret, frame = cap.read()
if not ret:
break
height, width, _ = frame.shape
line_position = width // 2 # Vertical line in the middle of the frame
# Reduce frame resolution for YOLO input
resized_frame = cv2.resize(frame, (320, 320))
blob = cv2.dnn.blobFromImage(resized_frame, 1/255.0, (320, 320), swapRB=True, crop=False)
net.setInput(blob)
# Get output layers names
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
# Forward pass
detections = net.forward(output_layers)
current_centroids = []
# Minimum confidence to consider as detection
confidence_threshold = 0.5 # Lowered threshold for higher sensitivity
nms_threshold = 0.4 # Non-Maximum Suppression threshold
boxes = []
confidences = []
class_ids = []
# Iterate over each detected object
for detection in detections:
for obj in detection:
scores = obj[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if class_id == 0 and confidence > confidence_threshold: # Class ID 0 is for person in YOLO
box = obj[0:4] * np.array([width, height, width, height])
(x_center, y_center, w, h) = box.astype("int")
x1 = int(x_center - w / 2)
y1 = int(y_center - h / 2)
x2 = x1 + w
y2 = y1 + h
boxes.append([x1, y1, x2 - x1, y2 - y1])
confidences.append(float(confidence))
class_ids.append(class_id)
# Apply Non-Maximum Suppression (NMS) to eliminate redundant overlapping boxes with lower confidences
indices = cv2.dnn.NMSBoxes(boxes, confidences, confidence_threshold, nms_threshold)
if len(indices) > 0:
for i in indices.flatten():
box = boxes[i]
x1, y1, w, h = box
x2 = x1 + w
y2 = y1 + h
# Draw bounding box
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
# Put label and confidence score
label = f'Person {confidences[i]:.2f}'
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# Calculate centroid of the bounding box
centroid = ((x1 + x2) // 2, (y1 + y2) // 2)
current_centroids.append(centroid)
# Update tracking information (assuming similar tracking logic as before)
updated_tracked_objects = {}
for centroid in current_centroids:
min_distance = float("inf")
assigned_id = None
# Find the closest tracked object
for person_id, data in tracked_objects.items():
prev_centroid = data['centroid']
distance = np.linalg.norm(np.array(centroid) - np.array(prev_centroid))
if distance < min_distance:
min_distance = distance
assigned_id = person_id
if min_distance > 50: # If no close tracked object is found, assign a new ID
assigned_id = next_person_id
next_person_id += 1
updated_tracked_objects[assigned_id] = {'centroid': centroid, 'counted': tracked_objects.get(assigned_id, {'counted': False})['counted']}
# Check if the object crosses the line
if assigned_id in tracked_objects:
prev_centroid = tracked_objects[assigned_id]['centroid']
if not tracked_objects[assigned_id]['counted'] and is_crossing_line(prev_centroid, centroid, line_position):
if prev_centroid[0] < line_position < centroid[0]:
people_in += 1
elif prev_centroid[0] > line_position > centroid[0]:
people_out += 1
updated_tracked_objects[assigned_id]['counted'] = True
tracked_objects = updated_tracked_objects
# Draw the vertical line
cv2.line(frame, (line_position, 0), (line_position, height), (0, 0, 255), 2)
# Show counts
total_people = people_in - people_out
cv2.putText(frame, f'IN: {people_in}', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
cv2.putText(frame, f'OUT: {people_out}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
cv2.putText(frame, f'Total: {total_people}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
# Send data to ESP32 only if there's a change in people_in or people_out
if people_in != last_sent_people_in or people_out != last_sent_people_out:
send_to_esp32(people_in, people_out, total_people)
last_sent_people_in = people_in
last_sent_people_out = people_out
# Show the frame with detections and counts
cv2.imshow('YOLOv3-tiny Detection', frame)
# Press 'q' to quit
if cv2.waitKey(1) == ord('q'):
break
# Ensure consistent frame rate
end_time = time.time()
elapsed_time = end_time - start_time
wait_time = max(1.0 / 15 - elapsed_time, 0) # Ensure consistent 15 FPS
time.sleep(wait_time)
cap.release()
cv2.destroyAllWindows()