一、检测图片
Python
import cv2
from ultralytics import YOLO
import torch
model_path = 'object_detection/best.pt' # Change this to your YOLOv8 model's path
image_path = 'object_detection/32.jpg' # Change this to your video's path
# Load the trained YOLOv8 model
model = YOLO(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device: %s" % device)
model.to(device)
# Process video frames
image = cv2.imread(image_path)
width, height, _ = image.shape
new_shape = [32*int(height/128), 32*int(width/128)]
image = cv2.resize(image, new_shape)
with torch.no_grad():
results = model.predict(image)
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = box.xyxy[0]
# Draw the bounding box on the BGR frame
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
# Add a label above the box
cv2.putText(image, result.names[int(box.cls)], (int(x1) - 30, int(y1) + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
cv2.imshow('Video', image)
cv2.waitKey(0)
cv2.destroyAllWindows()
二、检测视频
Python
import cv2
from ultralytics import YOLO
import torch
model_path = 'object_detection/best.pt' # Change this to your YOLOv8 model's path
video_path = 'object_detection/物块.mp4' # Change this to your video's path
# Load the trained YOLOv8 model
model = YOLO(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device: %s" % device)
model.to(device)
batch_size = 8
frames_rgb = []
frames = []
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Error: Could not open video.")
exit()
# Process video frames
while True:
ret, frame = cap.read()
if not ret:
print("Finished processing video.")
break
width, height, _ = frame.shape
new_shape = [32*int(height/64), 32*int(width/64)]
frame = cv2.resize(frame, new_shape)
frames.append(frame)
# YOLOv8 expects RGB images
if len(frames) == batch_size:
with torch.no_grad():
results = model.predict(frames)
# Process each detection
for i, result in enumerate(results):
for box in result.boxes:
print(box.conf)
if float(box.conf) > 0.9:
x1, y1, x2, y2 = box.xyxy[0]
# Draw the bounding box on the BGR frame
cv2.rectangle(frames[i], (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
# Add a label above the box
cv2.putText(frames[i], result.names[int(box.cls)], (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
cv2.imshow('Video', frames[i])
if cv2.waitKey(1) & 0xFF == ord('q'):
cap.release()
cv2.destroyAllWindows()
exit()
frames.clear()
frames_rgb.clear()
cap.release()
cv2.destroyAllWindows()
使用了sahi的视频检测
import argparse
import sys
import cv2
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction
import imageio
import numpy as np
def run(weights="yolov8n.pt", source="test.mp4", view_img=False):
"""
Run object detection on a video using YOLOv8 and SAHI.
Args:
weights (str): Model weights path.
source (str): Video file path.
view_img (bool): Show results.
"""
yolov8_model_path = weights
detection_model = AutoDetectionModel.from_pretrained(
model_type="yolov8", model_path=yolov8_model_path, confidence_threshold=0.3, device="cuda:0"
)
videocapture = cv2.VideoCapture(0)
new_shape = 32 * int(videocapture.get(3) / 64), 32 * int(videocapture.get(4) / 64)
writer = imageio.get_writer("object_detection/object_detection.mp4", fps=1 / 0.025)
while videocapture.isOpened():
success, frame = videocapture.read()
if not success:
break
frame = cv2.resize(frame, new_shape)
image = frame.copy()
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = get_sliced_prediction(
frame, detection_model, slice_height=512, slice_width=512, overlap_height_ratio=0.2, overlap_width_ratio=0.2
)
object_prediction_list = results.object_prediction_list
boxes_list = []
clss_list = []
for ind, _ in enumerate(object_prediction_list):
print(object_prediction_list[ind].score.value)
if float(object_prediction_list[ind].score.value) > 0.85:
boxes = (
object_prediction_list[ind].bbox.minx,
object_prediction_list[ind].bbox.miny,
object_prediction_list[ind].bbox.maxx,
object_prediction_list[ind].bbox.maxy,
)
clss = object_prediction_list[ind].category.name
boxes_list.append(boxes)
clss_list.append(clss)
for box, cls in zip(boxes_list, clss_list):
x1, y1, x2, y2 = box
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (56, 56, 255), 2)
label = str(cls)
t_size = cv2.getTextSize(label, 0, fontScale=0.6, thickness=1)[0]
cv2.rectangle(
image, (int(x1), int(y1) - t_size[1] - 3), (int(x1) + t_size[0], int(y1) + 3), (56, 56, 255), -1
)
cv2.putText(
image, label, (int(x1), int(y1) - 2), 0, 0.6, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA
)
if view_img:
cv2.imshow("result", image)
frame = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
writer.append_data((np.asarray(frame)).astype(np.uint8))
if cv2.waitKey(1) == ord("q"):
videocapture.release()
cv2.destroyAllWindows()
sys.exit()
writer.close()
def parse_opt():
"""Parse command line arguments."""
parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default="object_detection/best.pt", help="initial weights path")
parser.add_argument("--source", type=str, default="object_detection/物块.mp4", help="video file path")
parser.add_argument("--view-img", type=bool, default=True, help="show results")
return parser.parse_args()
def main(options):
"""Main function."""
run(**vars(options))
if __name__ == "__main__":
opt = parse_opt()
main(opt)