rawalkhirodkar commited on
Commit
cc0c583
1 Parent(s): 64ef3c8
Files changed (1) hide show
  1. app.py +16 -9
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
- keypoints[name] = (float(x), float(y), float(conf))
 
 
 
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 * bbox_width / 192) + x1
207
- y_coord = int(y * bbox_height / 256) + y1
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] * bbox_width / 192) + x1
221
- y1_coord = int(pt1[1] * bbox_height / 256) + y1
222
- x2_coord = int(pt2[0] * bbox_width / 192) + x1
223
- y2_coord = int(pt2[1] * bbox_height / 256) + y1
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)