92 lines
3.1 KiB
Python
92 lines
3.1 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
|
|
|
|
|
|
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'
|
|
}
|
|
|
|
# Load the pre-trained model (grid search object containing the best estimator)
|
|
model = joblib.load('svm_natural_images_modelUTC.pkl')
|
|
# If the saved object is a GridSearchCV, extract the best estimator.
|
|
if hasattr(model, 'best_estimator_'):
|
|
model = model.best_estimator_
|
|
|
|
# Detection parameters
|
|
windowSize = (64, 64)
|
|
hog_size = (64, 64)
|
|
stepSize = 8 # Step size for sliding window.
|
|
scale_factor = 1.01 # Scaling factor for image pyramid.
|
|
detection_threshold = 0.85 # Minimum probability to consider a detection valid.
|
|
|
|
# Load the test image (change the filename as needed)
|
|
image = cv2.imread("test3.jpg")
|
|
|
|
orig = image.copy()
|
|
detections = []
|
|
|
|
# Loop over the image pyramid
|
|
for resized in tuple(pyramid_gaussian(image, downscale=scale_factor, channel_axis=-1)):
|
|
# The ratio of the current resized image relative to the original image
|
|
scale_ratio = image.shape[1] / float(resized.shape[1])
|
|
|
|
# Slide a window across the resized image
|
|
for (x, y, window) in sliding_window(resized, stepSize, windowSize):
|
|
# Ensure the window is of the desired size
|
|
if window.shape[0] != windowSize[1] or window.shape[1] != windowSize[0]:
|
|
continue
|
|
|
|
if window.dtype != np.uint8:
|
|
window = (window * 255).astype(np.uint8)
|
|
|
|
|
|
gray = cv2.cvtColor(window, cv2.COLOR_BGR2GRAY)
|
|
gray = cv2.resize(gray, hog_size, interpolation=cv2.INTER_AREA)
|
|
|
|
features = hog(gray, **hog_params)
|
|
features = features.reshape(1, -1)
|
|
|
|
# Predict the class and probability for this window
|
|
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)
|
|
detections.append((startX, startY, endX, endY))
|
|
|
|
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()
|