init
This commit is contained in:
+130
@@ -0,0 +1,130 @@
|
||||
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):
|
||||
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]])
|
||||
|
||||
|
||||
# HOG parameters (should match those used during training)
|
||||
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) # Must match the window size used for training HOG features.
|
||||
stepSize = 8 # Step size for sliding window.
|
||||
scale_factor = 1.1 # Scaling factor for image pyramid.
|
||||
detection_threshold = 0.9 # Minimum probability to consider a detection valid.
|
||||
|
||||
# Load the pre-trained model (grid search object containing the best estimator)
|
||||
model = joblib.load('svm_natural_images_model.pkl')
|
||||
# If the saved object is a GridSearchCV, extract the best estimator.
|
||||
if hasattr(model, 'best_estimator_'):
|
||||
model = model.best_estimator_
|
||||
|
||||
|
||||
# This initializer will be run in each worker process to set globals.
|
||||
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
|
||||
|
||||
|
||||
def process_window(args):
|
||||
# Unpack arguments
|
||||
window, x, y, scale_ratio = args
|
||||
# Ensure proper datatype
|
||||
if window.dtype != np.uint8:
|
||||
window = (window * 255).astype(np.uint8)
|
||||
# Convert to grayscale and resize for HOG extraction
|
||||
gray = cv2.cvtColor(window, cv2.COLOR_BGR2GRAY)
|
||||
gray = cv2.resize(gray, hogSize)
|
||||
# Compute HOG features
|
||||
features = hog(gray, **hog_params)
|
||||
features = features.reshape(1, -1)
|
||||
# Make predictions
|
||||
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
|
||||
# Return bounding box if detection meets the criteria
|
||||
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
|
||||
|
||||
|
||||
async def async_detection(image):
|
||||
detections = []
|
||||
loop = asyncio.get_running_loop()
|
||||
tasks = []
|
||||
# Set up the ProcessPoolExecutor with max 10 processes and initializer
|
||||
with concurrent.futures.ProcessPoolExecutor(
|
||||
max_workers=20,
|
||||
initializer=init_worker,
|
||||
initargs=(model, hog_params, hogSize, detection_threshold, windowSize)
|
||||
) as executor:
|
||||
for resized in pyramid_gaussian(image, downscale=scale_factor, channel_axis=-1):
|
||||
scale_ratio = image.shape[1] / float(resized.shape[1])
|
||||
for (x, y, window) in sliding_window(resized, stepSize, windowSize):
|
||||
if window.shape[0] != windowSize[1] or window.shape[1] != windowSize[0]:
|
||||
continue
|
||||
args = (window, x, y, scale_ratio)
|
||||
tasks.append(loop.run_in_executor(executor, process_window, args))
|
||||
results = await asyncio.gather(*tasks)
|
||||
for result in results:
|
||||
if result is not None:
|
||||
detections.append(result)
|
||||
return detections
|
||||
|
||||
|
||||
async def main():
|
||||
# Load the test image (update the filename as needed)
|
||||
image = cv2.imread("test.jpg")
|
||||
if image is None:
|
||||
print("Failed to load image!")
|
||||
return
|
||||
|
||||
orig = image.copy()
|
||||
# Run the asynchronous detection
|
||||
detections = await async_detection(image)
|
||||
|
||||
# Apply non-maxima suppression to reduce overlapping boxes
|
||||
if len(detections) > 0:
|
||||
detections = np.array(detections)
|
||||
pick = non_max_suppression(detections, probs=None, overlapThresh=0.3)
|
||||
# Draw the final bounding boxes on the original image.
|
||||
for (xA, yA, xB, yB) in pick:
|
||||
cv2.rectangle(orig, (xA, yA), (xB, yB), (0, 255, 0), 2)
|
||||
|
||||
# Display the detections
|
||||
cv2.imshow("Detections", orig)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user