WideMax commited on
Commit
29b6ebb
1 Parent(s): 0241ca9

Update app.py

Browse files

To work correctly, you need to properly crop the image to the face, and resize it to a resolution of 1024. Both these actions are performed by align_face function, the method will work better with it.

Files changed (1) hide show
  1. app.py +5 -8
app.py CHANGED
@@ -32,20 +32,17 @@ os.chdir("..")
32
  shutil.rmtree("HairFastGAN")
33
 
34
  from hair_swap import HairFast, get_parser
 
35
 
36
  hair_fast = HairFast(get_parser().parse_args([]))
37
 
38
  def resize(image_path):
39
  img = Image.open(image_path)
40
- square_size = 1024
41
-
42
- left = (img.width - square_size) / 2
43
- top = (img.height - square_size) / 2
44
- right = (img.width + square_size) / 2
45
- bottom = (img.height + square_size) / 2
46
 
47
- img_cropped = img.crop((left, top, right, bottom))
48
- return img_cropped
49
 
50
  @spaces.GPU
51
  def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True)):
 
32
  shutil.rmtree("HairFastGAN")
33
 
34
  from hair_swap import HairFast, get_parser
35
+ from utils.shape_predictor import align_face
36
 
37
  hair_fast = HairFast(get_parser().parse_args([]))
38
 
39
  def resize(image_path):
40
  img = Image.open(image_path)
41
+
42
+ if img.size != (1024, 1024):
43
+ img = align_face(img, return_tensors=False)[0]
 
 
 
44
 
45
+ return img
 
46
 
47
  @spaces.GPU
48
  def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True)):