Don't include gradio

.gitignore

@@ -176,4 +176,7 @@ pyrightconfig.json
 # VSCode
 .vscode/
 
+# Gradio
+.gradio/
+
 # End of https://www.toptal.com/developers/gitignore/api/python

app.py

@@ -2,6 +2,10 @@ import gradio as gr
 import zombie
 from huggingface_hub import hf_hub_download
 import onnxruntime as ort
+import numpy as np
+from PIL import Image
+from faceparsing import get_face_mask
+
 # import torch
 # from your_pix2pixhd_code import YourPix2PixHDModel, load_image, tensor2im # Adapt these imports
 
@@ -30,27 +34,42 @@ ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider
 
 #     # return output_image
 
-def predict(input_image):
-    zombie_image = zombie.transition_onnx(input_image,ort_session)
-    
-    if zombie_image is None:
-        return "No face found"
-    
-    return zombie_image
+def predict(input_image, mode):    
+    if mode == "Classic":
+        # Use the transition_onnx function for side-by-side comparison
+        zombie_image = zombie.transition_onnx(input_image, ort_session)
+        if zombie_image is None:
+            return "No face found"
+        return zombie_image
+    elif mode == "In Place":
+        # Use the make_faces_zombie_from_array function for in-place transformation
+        #zombie_image = zombie.make_faces_zombie_from_array(im_array, None, ort_session)
+        #if zombie_image is None:
+        #    return "No face found"
+        #return zombie_image
+        face_mask = get_face_mask(input_image)
+        return face_mask
+
+    else:
+        return "Invalid mode selected"
 
 # --- 3. Create the Gradio Interface ---
-title = "pix2pixHD Image-to-Image Translation"
+title = "Make Me A Zombie"
 description = "Upload an image to see the pix2pixHD model in action."
-article = "<p style='text-align: center'>Model based on the <a href='https://github.com/NVIDIA/pix2pixHD' target='_blank'>pix2pixHD repository</a>.</p>"
+article = """<p style='text-align: center'>Model based on the <a href='https://github.com/NVIDIA/pix2pixHD' target='_blank'>pix2pixHD repository</a>.  
+More details at <a href='https://makemeazombie.com' target='_blank'>makemeazombie.com</a>.</p>"""
 
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Input Image"),
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Dropdown(choices=["Classic", "In Place"], value="Classic", label="Mode")
+    ],
     outputs=gr.Image(type="pil", label="Output Image"),
     title=title,
     description=description,
     article=article,
 )
 
-demo.launch()
-#demo.launch(debug=True)
+#demo.launch()
+demo.launch(debug=True)

faceparsing.py

@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+
+from PIL import Image
+import matplotlib.pyplot as plt
+
+# convenience expression for automatically determining device
+device = (
+    "cuda"
+    # Device for NVIDIA or AMD GPUs
+    if torch.cuda.is_available()
+    else "mps"
+    # Device for Apple Silicon (Metal Performance Shaders)
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+# load models
+image_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
+model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
+model.to(device)
+
+def get_face_mask(image):
+    # run inference on image
+    inputs = image_processor(images=image, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+    logits = outputs.logits  # shape (batch_size, num_labels, ~height/4, ~width/4)
+
+    # resize output to match input image dimensions
+    upsampled_logits = nn.functional.interpolate(logits,
+                    size=image.size[::-1], # H x W
+                    mode='bilinear',
+                    align_corners=False)
+
+    # get label masks
+    labels = upsampled_logits.argmax(dim=1)[0]
+
+    # move to CPU to visualize in matplotlib
+    labels_viz = labels.cpu().numpy()
+
+    #Map to something more colorful.  Use a color map to map the labels to a color.
+    #Create a color map for colors 0 through 18
+    color_map = plt.get_cmap('tab20')
+    #Map the labels to a color
+    colors = color_map(labels_viz)
+
+    #Convert to PIL Image
+    colors_pil = Image.fromarray((colors * 255).astype(np.uint8))
+
+
+    return labels_viz
+
+

requirements.txt

@@ -2,4 +2,6 @@ gradio
 onnxruntime-gpu
 opencv-python
 numpy
-mediapipe
+mediapipe
+torch
+transformers