HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add TrackZone test in test_solutions.py (#18411)

Browse source 
Signed-off-by: Muhammad Rizwan Munawar <muhammadrizwanmunawar123@gmail.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-12-30 19:53:39 +05:00 • committed by GitHub
parent 0e465175fd
commit 114bb0bdb6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 14 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

21
tests/test_solutions.py
View file

@ -14,46 +14,53 @@ POSE_VIDEO = "solution_ci_pose_demo.mp4"
@pytest.mark.slow @pytest.mark.slow
def test_major_solutions(): def test_major_solutions():
"""Test the object counting, heatmap, speed estimation and queue management solution.""" """Test the object counting, heatmap, speed estimation, trackzone and queue management solution."""
safe_download(url=f"{ASSETS_URL}/{DEMO_VIDEO}", dir=TMP) safe_download(url=f"{ASSETS_URL}/{DEMO_VIDEO}", dir=TMP)
cap = cv2.VideoCapture(str(TMP / DEMO_VIDEO)) cap = cv2.VideoCapture(str(TMP / DEMO_VIDEO))
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
region_points = [(20, 400), (1080, 400), (1080, 360), (20, 360)] region_points = [(20, 400), (1080, 400), (1080, 360), (20, 360)]
counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False) # Test object counter counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False) # Test object counter
heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False) # Test heatmaps heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False) # Test heatmaps
heatmap_count = solutions.Heatmap(
colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False, region=region_points
) # Test heatmaps with object counting
speed = solutions.SpeedEstimator(region=region_points, model="yolo11n.pt", show=False) # Test queue manager speed = solutions.SpeedEstimator(region=region_points, model="yolo11n.pt", show=False) # Test queue manager
queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False) # Test speed estimation queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False) # Test speed estimation
line_analytics = solutions.Analytics(analytics_type="line", model="yolo11n.pt", show=False) # line analytics line_analytics = solutions.Analytics(analytics_type="line", model="yolo11n.pt", show=False) # line analytics
pie_analytics = solutions.Analytics(analytics_type="pie", model="yolo11n.pt", show=False) # line analytics pie_analytics = solutions.Analytics(analytics_type="pie", model="yolo11n.pt", show=False) # line analytics
bar_analytics = solutions.Analytics(analytics_type="bar", model="yolo11n.pt", show=False) # line analytics bar_analytics = solutions.Analytics(analytics_type="bar", model="yolo11n.pt", show=False) # line analytics
area_analytics = solutions.Analytics(analytics_type="area", model="yolo11n.pt", show=False) # line analytics area_analytics = solutions.Analytics(analytics_type="area", model="yolo11n.pt", show=False) # line analytics
trackzone = solutions.TrackZone(region=region_points, model="yolo11n.pt", show=False) # Test trackzone
frame_count = 0 # Required for analytics frame_count = 0 # Required for analytics
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if not success: if not success:
break break
frame_count += 1
original_im0 = im0.copy() original_im0 = im0.copy()
_ = counter.count(original_im0.copy()) _ = counter.count(original_im0.copy())
_ = heatmap.generate_heatmap(original_im0.copy()) _ = heatmap.generate_heatmap(original_im0.copy())
_ = heatmap_count.generate_heatmap(original_im0.copy())
_ = speed.estimate_speed(original_im0.copy()) _ = speed.estimate_speed(original_im0.copy())
_ = queue.process_queue(original_im0.copy()) _ = queue.process_queue(original_im0.copy())
_ = line_analytics.process_data(original_im0.copy(), frame_count) _ = line_analytics.process_data(original_im0.copy(), frame_count)
_ = pie_analytics.process_data(original_im0.copy(), frame_count) _ = pie_analytics.process_data(original_im0.copy(), frame_count)
_ = bar_analytics.process_data(original_im0.copy(), frame_count) _ = bar_analytics.process_data(original_im0.copy(), frame_count)
_ = area_analytics.process_data(original_im0.copy(), frame_count) _ = area_analytics.process_data(original_im0.copy(), frame_count)
_ = trackzone.trackzone(original_im0.copy())
cap.release() cap.release()
# Test workouts monitoring # Test workouts monitoring
safe_download(url=f"{ASSETS_URL}/{POSE_VIDEO}", dir=TMP) safe_download(url=f"{ASSETS_URL}/{POSE_VIDEO}", dir=TMP)
cap1 = cv2.VideoCapture(str(TMP / POSE_VIDEO)) cap = cv2.VideoCapture(str(TMP / POSE_VIDEO))
assert cap1.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
gym = solutions.AIGym(line_width=2, kpts=[5, 11, 13], show=False) gym = solutions.AIGym(kpts=[5, 11, 13], show=False)
while cap1.isOpened(): while cap.isOpened():
success, im0 = cap1.read() success, im0 = cap.read()
if not success: if not success:
break break
_ = gym.monitor(im0) _ = gym.monitor(im0)
cap1.release() cap.release()
@pytest.mark.slow @pytest.mark.slow