Skip to content

Commit 4f23c91

Browse files
author
Hannah DeFazio
committed
Rescue some utils from deleted berkeley code
1 parent f80c0c1 commit 4f23c91

File tree

2 files changed

+231
-1
lines changed

2 files changed

+231
-1
lines changed
+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import os
2+
import math
3+
import ast
4+
import warnings
5+
6+
import ubelt as ub
7+
8+
9+
def load_hl_hand_bboxes(extracted_dir):
10+
fn = extracted_dir + "/_hand_pose_2d_data.json"
11+
12+
if not os.path.exists(fn):
13+
warnings.warn(f"{fn} does not exist, ignoring")
14+
return {}
15+
16+
with open(fn, "r") as f:
17+
hands = ast.literal_eval(f.read())
18+
19+
if hands == {} or hands == []:
20+
warnings.warn(f"hands data in {fn} is empty!")
21+
22+
all_hand_pose_2d = {}
23+
for hand_info in hands:
24+
time_stamp = float(hand_info["time_sec"]) + (
25+
float(hand_info["time_nanosec"]) * 1e-9
26+
)
27+
if time_stamp not in all_hand_pose_2d.keys():
28+
all_hand_pose_2d[time_stamp] = []
29+
30+
hand = hand_info["hand"].lower()
31+
hand_label = f"hand ({hand})"
32+
33+
joints = {}
34+
for joint in hand_info["joint_poses"]:
35+
# if joint['clipped'] == 0:
36+
joints[joint["joint"]] = joint # 2d position
37+
if joints != {}:
38+
all_hand_pose_2d[time_stamp].append({"hand": hand_label, "joints": joints})
39+
40+
return all_hand_pose_2d
41+
42+
43+
def add_hl_hand_bbox(preds):
44+
for video_name, dets in preds.items():
45+
all_hand_pose_2d_image_space = None
46+
47+
for frame, det in dets.items():
48+
meta = preds[video_name][frame]["meta"]
49+
time_stamp = meta["time_stamp"]
50+
# <video_folder>/_extracted/images/<file_name>
51+
video_folder = meta["file_name"].split("/")[:-3]
52+
video_folder = video_folder.join("/")
53+
54+
if not all_hand_pose_2d_image_space:
55+
all_hand_pose_2d_image_space = load_hl_hand_bboxes(
56+
video_folder + "/_extracted"
57+
)
58+
59+
# Add HL hand bounding boxes if we have them
60+
all_hands = (
61+
all_hand_pose_2d_image_space[time_stamp]
62+
if time_stamp in all_hand_pose_2d_image_space.keys()
63+
else []
64+
)
65+
if all_hands != []:
66+
print("Adding hand bboxes from the hololens joints")
67+
for joints in all_hands:
68+
keys = list(joints["joints"].keys())
69+
hand_label = joints["hand"]
70+
71+
all_x_values = [joints["joints"][k]["projected"][0] for k in keys]
72+
all_y_values = [joints["joints"][k]["projected"][1] for k in keys]
73+
74+
hand_bbox = [
75+
min(all_x_values),
76+
min(all_y_values),
77+
max(all_x_values),
78+
max(all_y_values),
79+
] # tlbr
80+
81+
new_det = {
82+
"confidence_score": 1,
83+
"bbox": hand_bbox,
84+
}
85+
preds[video_name][frame][hand_label] = [new_det]
86+
87+
return preds
88+
89+
90+
def find_closest_hands(object_pair, detected_classes, preds):
91+
# Determine what the hand label is in the video, if any
92+
# Fixes case where hand label has distinguishing information
93+
# ex: hand(right) vs hand (left)
94+
95+
hand_labels = [h for h in detected_classes if "hand" in h.lower()]
96+
97+
if len(hand_labels) == 0:
98+
return None
99+
# TODO: Update for multiple hand outputs
100+
return hand_labels
101+
102+
# find what object we should be interacting with
103+
try:
104+
obj = [o for o in object_pair if "hand" not in o][
105+
0
106+
] # What to do if we don't have this???
107+
obj_bbox = preds[obj]["bbox"]
108+
w = abs(obj_bbox[2] - obj_bbox[0])
109+
h = abs(obj_bbox[1] - obj_bbox[3])
110+
obj_center = [obj_bbox[0] + (w / 2), obj_bbox[1] + (h / 2)]
111+
except:
112+
return None # TODO: temp???
113+
114+
# Determine if any of the hands are close enough to the object to
115+
# likely be an interaction
116+
min_dist = 180
117+
close_hands = []
118+
for i, hand_label in enumerate(hand_labels):
119+
hand_bbox = preds[hand_label]["bbox"]
120+
w = abs(hand_bbox[2] - hand_bbox[0])
121+
h = abs(hand_bbox[1] - hand_bbox[3])
122+
hand_center = [hand_bbox[0] + (w / 2), hand_bbox[1] + (h / 2)]
123+
dist = math.dist(obj_center, hand_center)
124+
125+
if dist <= min_dist:
126+
close_hands.append(hand_label)
127+
128+
hand_label = close_hands if len(close_hands) > 0 else None
129+
return hand_label

angel_system/data/common/kwcoco_utils.py

+102-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
sanitize_str,
3131
)
3232
from angel_system.data.common.load_data import Re_order
33-
from angel_system.berkeley.data.update_dets_utils import load_hl_hand_bboxes
33+
from angel_system.data.common.hl_hands_utils import load_hl_hand_bboxes
3434

3535

3636
def load_kwcoco(dset):
@@ -135,6 +135,107 @@ def preds_to_kwcoco(
135135
return dset
136136

137137

138+
def replace_compound_label(
139+
preds,
140+
obj,
141+
detected_classes,
142+
using_contact,
143+
obj_hand_contact_state=False,
144+
obj_obj_contact_state=False,
145+
):
146+
"""
147+
Check if some subset of obj is in `detected_classes`
148+
149+
Designed to catch and correct incorrect compound classes
150+
Ex: obj is "mug + filter cone + filter" but we detected "mug + filter cone"
151+
"""
152+
153+
compound_objs = [x.strip() for x in obj.split("+")]
154+
detected_objs = [[y.strip() for y in x.split("+")] for x in detected_classes]
155+
156+
replaced_classes = []
157+
for detected_class, objs in zip(detected_classes, detected_objs):
158+
for obj_ in objs:
159+
if obj_ in compound_objs:
160+
replaced_classes.append(detected_class)
161+
break
162+
163+
if replaced_classes == []:
164+
# Case 0, We didn't detect any of the objects
165+
replaced = None
166+
elif len(replaced_classes) == 1:
167+
# Case 1, we detected a subset of the compound as one detection
168+
replaced = replaced_classes[0]
169+
# print(f'replaced {replaced} with {obj}')
170+
171+
preds[obj] = preds.pop(replaced)
172+
173+
if using_contact:
174+
for i in range(len(preds[obj])):
175+
preds[obj][i]["obj_hand_contact_state"] = obj_hand_contact_state
176+
preds[obj][i]["obj_obj_contact_state"] = obj_obj_contact_state
177+
else:
178+
# Case 2, the compound was detected as separate boxes
179+
replaced = replaced_classes
180+
# print(f'Combining {replaced} detections into compound \"{obj}\"')
181+
182+
new_bbox = None
183+
new_conf = None
184+
new_obj_obj_conf = None
185+
new_obj_hand_conf = None
186+
for det_obj in replaced:
187+
assert len(preds[det_obj]) == 1
188+
bbox = preds[det_obj][0]["bbox"]
189+
conf = preds[det_obj][0]["confidence_score"]
190+
191+
if new_bbox is None:
192+
new_bbox = bbox
193+
else:
194+
# Find mix of bboxes
195+
# TODO: first double check these are close enough that it makes sense to combine?
196+
new_tlx, new_tly, new_brx, new_bry = new_bbox
197+
tlx, tly, brx, bry = bbox
198+
199+
new_bbox = [
200+
min(new_tlx, tlx),
201+
min(new_tly, tly),
202+
max(new_brx, brx),
203+
max(new_bry, bry),
204+
]
205+
206+
new_conf = conf if new_conf is None else (new_conf + conf) / 2 # average???
207+
if using_contact:
208+
obj_obj_conf = preds[det_obj][0]["obj_obj_contact_conf"]
209+
new_obj_obj_conf = (
210+
obj_obj_conf
211+
if new_obj_obj_conf is None
212+
else (new_obj_obj_conf + obj_obj_conf) / 2
213+
)
214+
obj_hand_conf = preds[det_obj][0]["obj_hand_contact_conf"]
215+
new_obj_hand_conf = (
216+
obj_hand_conf
217+
if new_obj_hand_conf is None
218+
else (new_obj_hand_conf + obj_hand_conf) / 2
219+
)
220+
221+
# remove old preds
222+
preds.pop(det_obj)
223+
224+
new_pred = {
225+
"confidence_score": new_conf,
226+
"bbox": new_bbox,
227+
}
228+
if using_contact:
229+
new_pred["obj_obj_contact_state"] = obj_obj_contact_state
230+
new_pred["obj_obj_contact_conf"] = new_obj_obj_conf
231+
new_pred["obj_hand_contact_state"] = obj_hand_contact_state
232+
new_pred["obj_hand_contact_conf"] = new_obj_hand_conf
233+
234+
preds[obj] = [new_pred]
235+
236+
return preds, replaced
237+
238+
138239
def add_hl_hands_to_kwcoco(dset, remove_existing=True, using_contact=True):
139240
"""Add bounding boxes for the hands based on the Hololen's joint positions
140241

0 commit comments

Comments
 (0)