Spaces:
Running
on
Zero
Running
on
Zero
rawalkhirodkar
commited on
Commit
•
cc0c583
1
Parent(s):
64ef3c8
Add
Browse files
app.py
CHANGED
@@ -143,7 +143,7 @@ class ImageProcessor:
|
|
143 |
cropped_img = self.crop_image(result_image, bbox)
|
144 |
input_tensor = self.transform(cropped_img).unsqueeze(0).to("cuda")
|
145 |
heatmaps = ModelManager.run_model(pose_model, input_tensor)
|
146 |
-
keypoints = self.heatmaps_to_keypoints(heatmaps[0].cpu().numpy())
|
147 |
all_keypoints.append(keypoints) # Collect keypoints
|
148 |
result_image = self.draw_keypoints(result_image, keypoints, bbox, kpt_threshold)
|
149 |
|
@@ -166,15 +166,22 @@ class ImageProcessor:
|
|
166 |
return crop
|
167 |
|
168 |
@staticmethod
|
169 |
-
def heatmaps_to_keypoints(heatmaps):
|
170 |
num_joints = heatmaps.shape[0] # Should be 308
|
171 |
keypoints = {}
|
|
|
|
|
|
|
|
|
172 |
for i, name in enumerate(GOLIATH_KEYPOINTS):
|
173 |
if i < num_joints:
|
174 |
heatmap = heatmaps[i]
|
175 |
y, x = np.unravel_index(np.argmax(heatmap), heatmap.shape)
|
176 |
conf = heatmap[y, x]
|
177 |
-
|
|
|
|
|
|
|
178 |
return keypoints
|
179 |
|
180 |
@staticmethod
|
@@ -203,8 +210,8 @@ class ImageProcessor:
|
|
203 |
# Draw keypoints
|
204 |
for i, (name, (x, y, conf)) in enumerate(keypoints.items()):
|
205 |
if conf > kpt_threshold and i < len(GOLIATH_KPTS_COLORS):
|
206 |
-
x_coord = int(x
|
207 |
-
y_coord = int(y
|
208 |
color = GOLIATH_KPTS_COLORS[i]
|
209 |
cv2.circle(image, (x_coord, y_coord), radius, color, -1)
|
210 |
|
@@ -217,10 +224,10 @@ class ImageProcessor:
|
|
217 |
pt1 = keypoints[pt1_name]
|
218 |
pt2 = keypoints[pt2_name]
|
219 |
if pt1[2] > kpt_threshold and pt2[2] > kpt_threshold:
|
220 |
-
x1_coord = int(pt1[0]
|
221 |
-
y1_coord = int(pt1[1]
|
222 |
-
x2_coord = int(pt2[0]
|
223 |
-
y2_coord = int(pt2[1]
|
224 |
cv2.line(image, (x1_coord, y1_coord), (x2_coord, y2_coord), color, thickness=thickness)
|
225 |
|
226 |
return Image.fromarray(image)
|
|
|
143 |
cropped_img = self.crop_image(result_image, bbox)
|
144 |
input_tensor = self.transform(cropped_img).unsqueeze(0).to("cuda")
|
145 |
heatmaps = ModelManager.run_model(pose_model, input_tensor)
|
146 |
+
keypoints = self.heatmaps_to_keypoints(heatmaps[0].cpu().numpy(), bbox)
|
147 |
all_keypoints.append(keypoints) # Collect keypoints
|
148 |
result_image = self.draw_keypoints(result_image, keypoints, bbox, kpt_threshold)
|
149 |
|
|
|
166 |
return crop
|
167 |
|
168 |
@staticmethod
|
169 |
+
def heatmaps_to_keypoints(heatmaps, bbox):
|
170 |
num_joints = heatmaps.shape[0] # Should be 308
|
171 |
keypoints = {}
|
172 |
+
x1, y1, x2, y2 = map(int, bbox[:4])
|
173 |
+
bbox_width = x2 - x1
|
174 |
+
bbox_height = y2 - y1
|
175 |
+
|
176 |
for i, name in enumerate(GOLIATH_KEYPOINTS):
|
177 |
if i < num_joints:
|
178 |
heatmap = heatmaps[i]
|
179 |
y, x = np.unravel_index(np.argmax(heatmap), heatmap.shape)
|
180 |
conf = heatmap[y, x]
|
181 |
+
# Convert coordinates to image frame
|
182 |
+
x_image = x * bbox_width / 192 + x1
|
183 |
+
y_image = y * bbox_height / 256 + y1
|
184 |
+
keypoints[name] = (float(x_image), float(y_image), float(conf))
|
185 |
return keypoints
|
186 |
|
187 |
@staticmethod
|
|
|
210 |
# Draw keypoints
|
211 |
for i, (name, (x, y, conf)) in enumerate(keypoints.items()):
|
212 |
if conf > kpt_threshold and i < len(GOLIATH_KPTS_COLORS):
|
213 |
+
x_coord = int(x)
|
214 |
+
y_coord = int(y)
|
215 |
color = GOLIATH_KPTS_COLORS[i]
|
216 |
cv2.circle(image, (x_coord, y_coord), radius, color, -1)
|
217 |
|
|
|
224 |
pt1 = keypoints[pt1_name]
|
225 |
pt2 = keypoints[pt2_name]
|
226 |
if pt1[2] > kpt_threshold and pt2[2] > kpt_threshold:
|
227 |
+
x1_coord = int(pt1[0])
|
228 |
+
y1_coord = int(pt1[1])
|
229 |
+
x2_coord = int(pt2[0])
|
230 |
+
y2_coord = int(pt2[1])
|
231 |
cv2.line(image, (x1_coord, y1_coord), (x2_coord, y2_coord), color, thickness=thickness)
|
232 |
|
233 |
return Image.fromarray(image)
|