Create app.py

app.py

@@ -0,0 +1,54 @@
+import streamlit as st
+import torch
+from torch import nn
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+
+def load_model():
+    """Load the Segformer model and processor."""
+    processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
+    model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    model.to(device)
+    return processor, model, device
+
+def process_image(image: Image.Image, processor, model, device):
+    """Run inference on the image and return the segmentation mask."""
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    upsampled_logits = nn.functional.interpolate(
+        logits, size=image.size[::-1], mode="bilinear", align_corners=False
+    )
+    labels = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
+    return labels
+
+def visualize_segmentation(labels: np.ndarray):
+    """Visualize segmentation mask using Matplotlib."""
+    fig, ax = plt.subplots()
+    ax.imshow(labels, cmap="jet", alpha=0.7)
+    ax.axis("off")
+    st.pyplot(fig)
+
+# Streamlit UI
+st.title("Face Parsing using Segformer")
+st.write("Upload an image to perform semantic segmentation on faces.")
+
+# Load model
+processor, model, device = load_model()
+
+# File uploader
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+    
+    # Process image
+    with st.spinner("Processing..."):
+        labels = process_image(image, processor, model, device)
+        
+        # Display result
+        visualize_segmentation(labels)