xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1916 from AlexBodner/fix/docs-vitpose 

Corrected vitpose docstring
This commit is contained in:
Soumik Rakshit 2025-07-28 18:38:22 +05:30 committed by GitHub
parent 8873da4a30 9a0483446a
commit 8bbe80e683
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
supervision/keypoint/core.py
View File

@ -545,9 +545,9 @@ class KeyPoints:
DETECTION_MODEL_ID = "PekingU/rtdetr_r50vd_coco_o365"
detection_processor = AutoProcessor.from_pretrained(DETECTION_MODEL_ID, use_fast=True)
detection_model = RTDetrForObjectDetection.from_pretrained(DETECTION_MODEL_ID, device_map=DEVICE)
detection_model = RTDetrForObjectDetection.from_pretrained(DETECTION_MODEL_ID, device_map=device)
inputs = detection_processor(images=frame, return_tensors="pt").to(DEVICE)
inputs = detection_processor(images=frame, return_tensors="pt").to(device)
with torch.no_grad():
outputs = detection_model(**inputs)
@ -563,9 +563,9 @@ class KeyPoints:
pose_estimation_processor = AutoProcessor.from_pretrained(POSE_ESTIMATION_MODEL_ID)
pose_estimation_model = VitPoseForPoseEstimation.from_pretrained(
POSE_ESTIMATION_MODEL_ID, device_map=DEVICE)
POSE_ESTIMATION_MODEL_ID, device_map=device)
inputs = pose_estimation_processor(frame, boxes=[boxes], return_tensors="pt").to(DEVICE)
inputs = pose_estimation_processor(frame, boxes=[boxes], return_tensors="pt").to(device)
with torch.no_grad():
outputs = pose_estimation_model(**inputs)