RamAnanth1 commited on
Commit
8b3abf5
1 Parent(s): f016e80

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +4 -4
model.py CHANGED
@@ -162,8 +162,6 @@ class Model:
162
  subprocess.run(shlex.split(f'wget {pose_checkpoint_file} -O models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'))
163
  pose_checkpoint = 'models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'
164
 
165
- det_cat_id = 1
166
- bbox_thr = 0.2
167
  ## detector
168
  det_config_mmcv = mmcv.Config.fromfile(det_config)
169
  self.det_model = init_detector(det_config_mmcv, det_checkpoint, device=device)
@@ -280,6 +278,8 @@ class Model:
280
  @torch.inference_mode()
281
  def process_pose(self, input_img, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):
282
  global current_base
 
 
283
  device = 'cuda'
284
  if base_model == 'sd-v1-4.ckpt':
285
  model = self.model
@@ -302,7 +302,7 @@ class Model:
302
 
303
  image = im.copy()
304
  im = img2tensor(im).unsqueeze(0)/255.
305
- mmdet_results = inference_detector(det_model, image)
306
  # keep the person class bounding boxes.
307
  person_results = process_mmdet_results(mmdet_results, det_cat_id)
308
 
@@ -313,7 +313,7 @@ class Model:
313
  # e.g. use ('backbone', ) to return backbone feature
314
  output_layer_names = None
315
  pose_results, returned_outputs = inference_top_down_pose_model(
316
- pose_model,
317
  image,
318
  person_results,
319
  bbox_thr=bbox_thr,
 
162
  subprocess.run(shlex.split(f'wget {pose_checkpoint_file} -O models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'))
163
  pose_checkpoint = 'models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'
164
 
 
 
165
  ## detector
166
  det_config_mmcv = mmcv.Config.fromfile(det_config)
167
  self.det_model = init_detector(det_config_mmcv, det_checkpoint, device=device)
 
278
  @torch.inference_mode()
279
  def process_pose(self, input_img, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):
280
  global current_base
281
+ det_cat_id = 1
282
+ bbox_thr = 0.2
283
  device = 'cuda'
284
  if base_model == 'sd-v1-4.ckpt':
285
  model = self.model
 
302
 
303
  image = im.copy()
304
  im = img2tensor(im).unsqueeze(0)/255.
305
+ mmdet_results = inference_detector(self.det_model, image)
306
  # keep the person class bounding boxes.
307
  person_results = process_mmdet_results(mmdet_results, det_cat_id)
308
 
 
313
  # e.g. use ('backbone', ) to return backbone feature
314
  output_layer_names = None
315
  pose_results, returned_outputs = inference_top_down_pose_model(
316
+ self.pose_model,
317
  image,
318
  person_results,
319
  bbox_thr=bbox_thr,