face parsing demo working locally

faceparsing.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+import numpy as np
 
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -38,17 +39,24 @@ def get_face_mask(image):
 
     # move to CPU to visualize in matplotlib
     labels_viz = labels.cpu().numpy()
+    
+    # Debug: print label statistics
+    print(f"Labels min: {labels_viz.min()}, max: {labels_viz.max()}, unique: {np.unique(labels_viz)}")
 
     #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)
+    #Map the labels to a color - normalize labels to 0-1 range for the colormap
+    # Face parsing models typically have 19 classes (0-18), so normalize by 18
+    normalized_labels = labels_viz.astype(np.float32) / 18.0
+    colors = color_map(normalized_labels)
 
-    #Convert to PIL Image
-    colors_pil = Image.fromarray((colors * 255).astype(np.uint8))
 
+    #Convert to PIL Image - take only RGB channels (drop alpha)
+    colors_rgb = colors[:, :, :3]  # Remove alpha channel
+    colors_pil = Image.fromarray((colors_rgb * 255).astype(np.uint8))
 
-    return labels_viz
+
+    return colors_pil