118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
import joblib
|
|
from skimage.feature import hog
|
|
import imutils
|
|
from imutils.object_detection import non_max_suppression
|
|
from skimage.transform import pyramid_gaussian
|
|
import asyncio
|
|
import concurrent.futures
|
|
|
|
def sliding_window(image, stepSize, windowSize):
|
|
# Standard (synchronous) sliding window generator (not used in async version)
|
|
for y in range(0, image.shape[0] - windowSize[1] + 1, stepSize):
|
|
for x in range(0, image.shape[1] - windowSize[0] + 1, stepSize):
|
|
yield (x, y, image[y:y + windowSize[1], x:x + windowSize[0]])
|
|
|
|
# Helper function: generate sliding windows for a given y-range
|
|
def generate_windows_for_y_range(image, stepSize, windowSize, y_start, y_end):
|
|
windows = []
|
|
# Ensure y_end does not exceed allowed range
|
|
for y in range(y_start, min(y_end, image.shape[0] - windowSize[1] + 1), stepSize):
|
|
for x in range(0, image.shape[1] - windowSize[0] + 1, stepSize):
|
|
window = image[y:y + windowSize[1], x:x + windowSize[0]]
|
|
# Only add if the window has the desired size
|
|
if window.shape[0] == windowSize[1] and window.shape[1] == windowSize[0]:
|
|
windows.append((x, y, window))
|
|
return windows
|
|
|
|
# Asynchronous sliding window generator using 5 workers
|
|
async def async_sliding_window(image, stepSize, windowSize, num_workers=5):
|
|
loop = asyncio.get_running_loop()
|
|
height = image.shape[0] - windowSize[1] + 1
|
|
segment_size = max(1, height // num_workers)
|
|
tasks = []
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
|
|
for i in range(0, height, segment_size):
|
|
y_start = i
|
|
y_end = i + segment_size
|
|
tasks.append(loop.run_in_executor(
|
|
executor,
|
|
generate_windows_for_y_range,
|
|
image, stepSize, windowSize, y_start, y_end
|
|
))
|
|
results = await asyncio.gather(*tasks)
|
|
# Flatten the list of lists into a single list of windows
|
|
windows = [win for sublist in results for win in sublist]
|
|
return windows
|
|
|
|
# HOG parameters (must match training parameters)
|
|
hog_params = {
|
|
'orientations': 9,
|
|
'pixels_per_cell': (8, 8),
|
|
'cells_per_block': (2, 2),
|
|
'block_norm': 'L2-Hys'
|
|
}
|
|
|
|
# Detection parameters
|
|
windowSize = (64, 64)
|
|
hogSize = (128, 128) # Size to which each window is resized for HOG extraction.
|
|
stepSize = 8 # Sliding window step size.
|
|
scale_factor = 1.1 # Pyramid downscale factor.
|
|
detection_threshold = 0.9 # Minimum probability to consider a detection valid.
|
|
|
|
# Load the pre-trained model
|
|
model = joblib.load('svm_natural_images_model.pkl')
|
|
if hasattr(model, 'best_estimator_'):
|
|
model = model.best_estimator_
|
|
|
|
# Initializer for each worker process
|
|
def init_worker(_model, _hog_params, _hogSize, _detection_threshold, _windowSize):
|
|
global model, hog_params, hogSize, detection_threshold, windowSize
|
|
model = _model
|
|
hog_params = _hog_params
|
|
hogSize = _hogSize
|
|
detection_threshold = _detection_threshold
|
|
windowSize = _windowSize
|
|
|
|
# Process a single window: extract features and run prediction
|
|
def process_window(args):
|
|
window, x, y, scale_ratio = args
|
|
if window.dtype != np.uint8:
|
|
window = (window * 255).astype(np.uint8)
|
|
# Convert to grayscale and resize for HOG
|
|
gray = cv2.cvtColor(window, cv2.COLOR_BGR2GRAY)
|
|
gray = cv2.resize(gray, hogSize)
|
|
features = hog(gray, **hog_params)
|
|
features = features.reshape(1, -1)
|
|
pred = model.predict(features)
|
|
prob = model.predict_proba(features)[0]
|
|
if 'person' in model.classes_:
|
|
person_index = list(model.classes_).index('person')
|
|
person_prob = prob[person_index]
|
|
else:
|
|
person_prob = 0
|
|
if pred[0] == 'person' and person_prob > detection_threshold:
|
|
startX = int(x * scale_ratio)
|
|
startY = int(y * scale_ratio)
|
|
endX = int((x + windowSize[0]) * scale_ratio)
|
|
endY = int((y + windowSize[1]) * scale_ratio)
|
|
return (startX, startY, endX, endY)
|
|
return None
|
|
|
|
# Main asynchronous detection function
|
|
async def async_detection(image):
|
|
detections = []
|
|
loop = asyncio.get_running_loop()
|
|
tasks = []
|
|
# ProcessPoolExecutor for heavy computation (up to 10 processes)
|
|
with concurrent.futures.ProcessPoolExecutor(
|
|
max_workers=10,
|
|
initializer=init_worker,
|
|
initargs=(model, hog_params, hogSize, detection_threshold, windowSize)
|
|
) as process_executor:
|
|
# Iterate over each level in the image pyramid
|
|
for resized in pyramid_gaussian(image, downscale=scale_factor, channel_axis=-1):
|
|
scale_ratio = image.shape[1] / float(resized.shape[1])
|
|
# Asynchronously generate sliding
|