Update swap.py

swap.py

@@ -3,29 +3,27 @@ import os
 
 import cv2
 import torch
-import Image
+import numpy as np  # Import numpy explicitly
+from PIL import Image  # Use PIL for image processing
 from insightface.app import FaceAnalysis
 import face_align
 
 faceAnalysis = FaceAnalysis(name='buffalo_l')
-faceAnalysis.prepare(ctx_id=0, det_size=(512, 512))
+faceAnalysis.prepare(ctx_id=-1, det_size=(512, 512)) #ctx_id=-1 for CPU
 
 from StyleTransferModel_128 import StyleTransferModel
+import gradio as gr
 
 def parse_arguments():
     parser = argparse.ArgumentParser(description='Process command line arguments')
     
-    parser.add_argument('--target', required=True, help='Target path')
-    parser.add_argument('--source', required=True, help='Source path')
-    parser.add_argument('--outputPath', required=True, help='Output path')
     parser.add_argument('--modelPath', required=True, help='Model path')
-    parser.add_argument('--no-paste-back', action='store_true', help='Disable pasting back the swapped face onto the original image')
     parser.add_argument('--resolution', type=int, default=128, help='Resolution')
 
     return parser.parse_args()
 
 def get_device():
-    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    return torch.device('cpu')  # Force CPU
 
 def load_model(model_path):
     device = get_device()
@@ -43,51 +41,78 @@ def swap_face(model, target_face, source_face_latent):
     with torch.no_grad():
         swapped_tensor = model(target_tensor, source_tensor)
     
-    swapped_face = Image.postprocess_face(swapped_tensor)
-    
+    swapped_face = postprocess_face(swapped_tensor) # Use PIL-based postprocess
+
     return swapped_face, swapped_tensor
 
 def create_target(target_image, resolution):
-    if isinstance(target_image, str):
-        target_image = cv2.imread(target_image)
+    target_face = faceAnalysis.get(np.array(target_image))[0] # Convert PIL to numpy
 
-    target_face = faceAnalysis.get(target_image)[0]
-    aligned_target_face, M = face_align.norm_crop2(target_image, target_face.kps, resolution)
-    target_face_blob = Image.getBlob(aligned_target_face, (resolution, resolution))
+    aligned_target_face, M = face_align.norm_crop2(np.array(target_image), target_face.kps, resolution) # Convert PIL to numpy
+    target_face_blob = getBlob(aligned_target_face, (resolution, resolution))
 
     return target_face_blob, M
 
-def create_source(source_img_path):
-    source_image = cv2.imread(source_img_path)
+def create_source(source_image):
+    source_face = faceAnalysis.get(np.array(source_image))[0] # Convert PIL to numpy
+    source_latent = getLatent(source_face)
 
-    source_face = faceAnalysis.get(source_image)[0]
+    return source_latent
 
-    source_latent = Image.getLatent(source_face)
 
-    return source_latent
+def postprocess_face(swapped_tensor):
+    swapped_tensor = swapped_tensor.cpu().numpy()
+    swapped_tensor = np.transpose(swapped_tensor, (0, 2, 3, 1))
+    swapped_tensor = (swapped_tensor * 255).astype(np.uint8)
+    swapped_face = Image.fromarray(swapped_tensor[0]) # Convert to PIL Image
+    return swapped_face
+
+def getBlob(aligned_face, size):
+    aligned_face = cv2.resize(aligned_face, size)
+    aligned_face = aligned_face / 255.0
+    aligned_face = np.transpose(aligned_face, (2, 0, 1))
+    aligned_face = np.expand_dims(aligned_face, axis=0)
+    aligned_face = torch.from_numpy(aligned_face).float()
+    return aligned_face
+
+def getLatent(source_face):
+    return source_face.embedding
 
-def main():
-    args = parse_arguments()
-    
-    # Access the arguments
-    target_image_path = args.target
-    source = args.source
-    output_path = args.outputPath
-    model_path = args.modelPath
 
-    model = load_model(model_path)
+def blend_swapped_image(swapped_face, target_img, M):
+    swapped_face = np.array(swapped_face) # PIL to numpy
+    swapped_face = cv2.warpAffine(swapped_face, M, (target_img.shape[1], target_img.shape[0]))
+    mask = np.ones_like(swapped_face) * 255
+    mask = cv2.warpAffine(mask, M, (target_img.shape[1], target_img.shape[0]))
 
-    target_img = cv2.imread(target_image_path)
-    target_face_blob, M = create_target(target_img, args.resolution)
-    source_latent = create_source(source)
+    target_img = np.array(target_img) # PIL to numpy
+    swapped_face = Image.blend(Image.fromarray(target_img), Image.fromarray(swapped_face), Image.fromarray(mask).convert("L"))
+
+    return np.array(swapped_face) # numpy to PIL
+
+
+def process_images(target_image, source_image, model_path):
+    args = parse_arguments()
+    args.modelPath = model_path
+    args.no_paste_back = False # or True, as you prefer
+    args.resolution = 128
+    model = load_model(args.modelPath)
+
+    target_face_blob, M = create_target(target_image, args.resolution)
+    source_latent = create_source(source_image)
     swapped_face, _ = swap_face(model, target_face_blob, source_latent)
 
-    if not args.no_paste_back:
-        swapped_face = Image.blend_swapped_image(swapped_face, target_img, M)
+    swapped_face = blend_swapped_image(swapped_face, target_image, M) # PIL images
+
+    return Image.fromarray(swapped_face) # Return PIL image
+
 
-    output_folder = os.path.dirname(output_path)
-    os.makedirs(output_folder, exist_ok=True)
-    cv2.imwrite(output_path, swapped_face)
+with gr.Blocks() as demo:
+    target_image = gr.Image(label="Target Image", type="pil") # Use PIL type
+    source_image = gr.Image(label="Source Image", type="pil") # Use PIL type
+    model_path = gr.Textbox(label="Model Path", value="path/to/your/model.pth") # Add model path input
+    output_image = gr.Image(label="Output Image", type="pil") # Use PIL type
+    btn = gr.Button("Swap Face")
+    btn.click(fn=process_images, inputs=[target_image, source_image, model_path], outputs=output_image)
 
-if __name__ == "__main__":
-    main()
+demo.launch() #no share = true for local running