Overlaying a heat map over your webcam with YoloV8 + Ultralytics
I have been messing around a lot with OpenCV and one of the cool things I’ve been able to get working is running tracking on live camera feeds. One of those things I’ve been messing with is heatmaps.
from ultralytics import YOLO from ultralytics.solutions import heatmap import cv2 model = YOLO("yolov8n.pt") #Initializes capture, 0 is default webcam cap = cv2.VideoCapture(0) assert cap.isOpened(), "Error reading video file" #video properties w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) # Video writer settings video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) # Init heatmap heatmap_obj = heatmap.Heatmap() #Set properties for heatmap heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA, imw=w, imh=h, view_img=True, shape="circle") #Processing loop fo while the capture is open while cap.isOpened(): success, im0 = cap.read() if not success: print("Video frame is empty or video processing has been successfully completed.") break #initializes Yolo tracking, persist is for consitency, and show=false removes default vizualization tracks = model.track(im0, persist=True, show=False) #generates heatmap on tracked object im0 = heatmap_obj.generate_heatmap(im0, tracks) #write the frame with heatmap over video video_writer.write(im0) # Display the frame with heatmap, commented out becasue with yolo it makes 2 displays #cv2.imshow('Heatmap', im0) #leave cam when pressed if cv2.waitKey(1) & 0xFF == ord('q'): # Break the loop with the 'q' key break #closing logic cap.release() video_writer.release() cv2.destroyAllWindows()
To get started you need to download the openCV and ultralytics libraries.
OpenCV is for the webcam. Ultralytics is for the heatmap.
#Initializes capture, 0 is default webcam cap = cv2.VideoCapture(0) assert cap.isOpened(), "Error reading video file" #video properties w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) # Video writer settings video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
This is the code for initializing the capture. Normally 0 is the default webcam for a computer, but you can sync it with connected cameras instead. Video writer settings is just for writing the end product to a video file.
# Init heatmap heatmap_obj = heatmap.Heatmap() #Set properties for heatmap heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA, imw=w, imh=h, view_img=True, shape="circle")
This is the code for initializing the heatmap object. in .set_args, I set the parameters I want for the heatmap. imw is the width of the frame, imh is the height of the frame, colormap is the type of colormap I’m using, and shape is pretty self explanatory.
#Processing loop fo while the capture is open while cap.isOpened(): success, im0 = cap.read() if not success: print("Video frame is empty or video processing has been successfully completed.") break #initializes Yolo tracking, persist is for consitency, and show=false removes default vizualization tracks = model.track(im0, persist=True, show=False) #generates heatmap on tracked object im0 = heatmap_obj.generate_heatmap(im0, tracks) #write the frame with heatmap over video video_writer.write(im0) # Display the frame with heatmap, commented out becasue with yolo it makes 2 displays #cv2.imshow('Heatmap', im0) #leave cam when pressed if cv2.waitKey(1) & 0xFF == ord('q'): # Break the loop with the 'q' key break
This is the loop for video processing. cap.isOpened means that while the video capture is open it’ll run everything within that capture.
#initializes Yolo tracking, persist is for consitency, and show=false removes default vizualization tracks = model.track(im0, persist=True, show=False) #generates heatmap on tracked object im0 = heatmap_obj.generate_heatmap(im0, tracks) #write the frame with heatmap over video video_writer.write(im0)
, tracks = the Yolo model initialized above. This is pretty much setting up how to track the heat sigs. im0 will generate it, and the video_writer.write will write the heatmap over the video.
the cv2.waitKey code sets it so if I press Q and hold for a second it’ll end the program.
But yea thats pretty much how the code works. Important link right here for ultralytics and their guides:
