Add face detection

app.py

@@ -6,9 +6,9 @@ import PIL
 from PIL import Image
 import numpy as np
 import os
+import cvlib as cv
 import uuid
 import torch
-from torch import autocast
 import cv2
 from matplotlib import pyplot as plt
 from torchvision import transforms
@@ -20,7 +20,7 @@ auth_token = os.environ.get("API_TOKEN") or True
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", dtype=torch.float16, revision="fp16", use_auth_token=auth_token).to(device)
+pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", dtype=torch.float32, revision="fp16", use_auth_token=auth_token).to(device)
 
 transform = transforms.Compose([
       transforms.ToTensor(),
@@ -38,11 +38,24 @@ def read_content(file_path: str) -> str:
 
 def predict(dict, prompt=""):
     init_image = dict["image"].convert("RGB").resize((512, 512))
-    mask = dict["mask"].convert("RGB").resize((512, 512))
-    output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)
+    _init_image = cv2.cvtColor(np.array(init_image), cv2.COLOR_RGB2BGR)
+    faces, confidences = cv.detect_face(_init_image)
+    cv2.imwrite('data/init_image.jpg',_init_image)
+    for (x,y,p,q) in faces:
+        cv2.rectangle(_init_image,(x,y),(p,q),(255,0,0),2)
+    cv2.imwrite('data/face_detected_image.jpg',_init_image)
+    
+    (x, y, x2, y2) = faces[0]
+
+    face_mask = np.zeros((512, 512))
+    face_mask[y:y2, x:x2] = 255
+    cv2.imwrite('data/face_mask.jpg',face_mask)
+    mask = Image.fromarray(face_mask).convert("RGB")
+
+    # mask = dict["mask"].convert("RGB").resize((512, 512))
+    output = pipe(prompt = prompt, image=init_image, mask_image=mask, guidance_scale=8) #7.5
     return output.images[0], gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
-
 css = '''
 .container {max-width: 1150px;margin: auto;padding-top: 1.5rem}
 #image_upload{min-height:400px}
@@ -92,7 +105,7 @@ with image_blocks as demo:
                     image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload source image here").style(height=400)
                     with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                         prompt = gr.Textbox(placeholder = 'Enter name here (what you want in place of what is erased)', show_label=False, elem_id="input-text")
-                        btn = gr.Button("Abracadabra!").style(
+                        btn = gr.Button("Generate!").style(
                             margin=False,
                             rounded=(False, True, True, False),
                             full_width=False,
@@ -125,4 +138,4 @@ with image_blocks as demo:
                 """
             )
 
-image_blocks.launch()
+image_blocks.launch(share=True)

requirements.txt

@@ -8,4 +8,6 @@ numpy
 matplotlib
 uuid
 opencv-python
-git+https://github.com/openai/CLIP.git
+tensorflow
+cvlib
+git+https://github.com/openai/CLIP.git