55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
import os
|
|
from ultralytics import YOLO
|
|
|
|
def main():
|
|
# 1. Paths
|
|
model_path = "person_detector_best.pt"
|
|
video_path = os.path.join(
|
|
"video_data",
|
|
"S3 - 55 40 WJ -7",
|
|
"download_ch18_20260406080008_20260406090008.mp4"
|
|
)
|
|
|
|
# 2. Validations
|
|
if not os.path.exists(model_path):
|
|
print(f"[!] Error: Model not found at '{model_path}'")
|
|
return
|
|
|
|
if not os.path.exists(video_path):
|
|
print(f"[!] Error: Video not found at '{video_path}'")
|
|
return
|
|
|
|
# 3. Load Model
|
|
print(f"[*] Loading model: {model_path}")
|
|
# Using the Ultralytics YOLO class
|
|
model = YOLO(model_path)
|
|
|
|
# 4. Run Inference
|
|
print(f"[*] Running inference on video: {video_path}")
|
|
print("[*] Press 'q' in the video window to stop inference early.")
|
|
|
|
# 'stream=True' is memory-efficient and processes the video frame by frame
|
|
# 'save=True' will save the annotated video to 'runs/detect/predict'
|
|
# 'show=True' will display the video as it processes
|
|
results = model.predict(
|
|
source=video_path,
|
|
conf=0.35, # Confidence threshold (adjust if needed)
|
|
save=True, # Save the output video
|
|
show=True, # Show video in real-time
|
|
stream=True, # Important for large 2GB video files
|
|
imgsz=1280,
|
|
)
|
|
|
|
# Since stream=True, we must iterate over the results generator to execute inference
|
|
try:
|
|
for _ in results:
|
|
pass
|
|
except KeyboardInterrupt:
|
|
print("\n[!] Inference interrupted by user.")
|
|
|
|
print("\n[*] Inference finished.")
|
|
print("[*] Check the 'runs/detect/predict' directory (or similar) for your annotated video.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|