spdraptor commited on
Commit
23df3d5
Β·
1 Parent(s): cab5c2f

add mask area option

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -122,7 +122,7 @@ pipe = TryonPipeline.from_pretrained(
122
  pipe.unet_encoder = UNet_Encoder
123
 
124
  @spaces.GPU
125
- def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed):
126
  device = "cuda"
127
 
128
  openpose_model.preprocessor.body_estimation.model.to(device)
@@ -150,7 +150,7 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
150
  if is_checked:
151
  keypoints = openpose_model(human_img.resize((384,512)))
152
  model_parse, _ = parsing_model(human_img.resize((384,512)))
153
- mask, mask_gray = get_mask_location('hd', "dresses", model_parse, keypoints)
154
  mask = mask.resize((768,1024))
155
  else:
156
  mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
@@ -269,6 +269,8 @@ with image_blocks as demo:
269
  is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
270
  with gr.Row():
271
  is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
 
 
272
 
273
  example = gr.Examples(
274
  inputs=imgs,
@@ -305,7 +307,7 @@ with image_blocks as demo:
305
 
306
 
307
 
308
- try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
309
 
310
 
311
 
 
122
  pipe.unet_encoder = UNet_Encoder
123
 
124
  @spaces.GPU
125
+ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,mask_area,denoise_steps,seed):
126
  device = "cuda"
127
 
128
  openpose_model.preprocessor.body_estimation.model.to(device)
 
150
  if is_checked:
151
  keypoints = openpose_model(human_img.resize((384,512)))
152
  model_parse, _ = parsing_model(human_img.resize((384,512)))
153
+ mask, mask_gray = get_mask_location('hd', mask_area, model_parse, keypoints)
154
  mask = mask.resize((768,1024))
155
  else:
156
  mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
 
269
  is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
270
  with gr.Row():
271
  is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
272
+ with gr.Row():
273
+ mask_area = gr.Radio(['upper_body', 'lower_body', 'dresses'], label="mask area", info="select what part to mask")
274
 
275
  example = gr.Examples(
276
  inputs=imgs,
 
307
 
308
 
309
 
310
+ try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop,mask_area, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
311
 
312
 
313