Rawal Khirodkar commited on
Commit
1bdad69
·
1 Parent(s): 7ca9191

Pose: per-person bbox-scaled radius+thickness; green bbox overlay

Browse files
Files changed (1) hide show
  1. app.py +25 -12
app.py CHANGED
@@ -185,18 +185,31 @@ def predict(image: Image.Image, size: str, kpt_thr: float):
185
  keypoints, scores = _estimate_pose(image_bgr, bboxes, model)
186
 
187
  meta = model.pose_metainfo
188
- vis_rgb = visualize_keypoints(
189
- image=image_rgb,
190
- keypoints=keypoints,
191
- keypoints_visible=[np.ones(len(s), dtype=bool) for s in scores],
192
- keypoint_scores=scores,
193
- radius=3,
194
- thickness=1,
195
- kpt_thr=kpt_thr,
196
- skeleton=meta["skeleton_links"],
197
- kpt_color=meta["keypoint_colors"],
198
- link_color=meta["skeleton_link_colors"],
199
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
200
  vis_pil = Image.fromarray(vis_rgb)
201
 
202
  instances = [
 
185
  keypoints, scores = _estimate_pose(image_bgr, bboxes, model)
186
 
187
  meta = model.pose_metainfo
188
+ vis_rgb = image_rgb.copy()
189
+ for bbox, kpts, scr in zip(bboxes, keypoints, scores):
190
+ x1, y1, x2, y2 = map(int, bbox[:4])
191
+ # Scale render sizes with bbox area (per-person), matching v1's recipe.
192
+ bbox_diag = float(np.sqrt(max(1, x2 - x1) * max(1, y2 - y1)))
193
+ radius = max(1, int(bbox_diag * 0.006))
194
+ thickness = max(1, int(bbox_diag * 0.006))
195
+ bbox_thickness = max(1, thickness // 4)
196
+
197
+ # Green bbox first (so skeleton draws on top).
198
+ cv2.rectangle(vis_rgb, (x1, y1), (x2, y2), (0, 255, 0), bbox_thickness)
199
+
200
+ # Skeleton + keypoints for this person only.
201
+ vis_rgb = visualize_keypoints(
202
+ image=vis_rgb,
203
+ keypoints=[kpts],
204
+ keypoints_visible=[np.ones(len(scr), dtype=bool)],
205
+ keypoint_scores=[scr],
206
+ radius=radius,
207
+ thickness=thickness,
208
+ kpt_thr=kpt_thr,
209
+ skeleton=meta["skeleton_links"],
210
+ kpt_color=meta["keypoint_colors"],
211
+ link_color=meta["skeleton_link_colors"],
212
+ )
213
  vis_pil = Image.fromarray(vis_rgb)
214
 
215
  instances = [