Upload AuraBloom.py

AuraBloom.py

@@ -0,0 +1,134 @@
+import streamlit as st
+from PIL import Image
+import numpy as np
+from joblib import load
+from skimage.transform import resize
+import torch
+import os
+import sys
+
+# Ensure to run these commands in your terminal first:
+# pip install git+https://github.com/FacePerceiver/facer.git@main
+# pip install timm
+# git clone https://github.com/FacePerceiver/facer.git
+
+# Set the path for the 'facer' module
+sys.path.append('facer')
+
+import facer
+
+# Load face parsing model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+face_parser = facer.face_parser('farl/lapa/448', device=device)
+
+# Define the monk scale colors
+monk_scale = {
+    'Class2': (243, 231, 219),  # f3e7db
+    'Class3': (247, 234, 208),  # f7ead0
+    'Class4': (234, 218, 186),  # eadaba
+    'Class5': (215, 189, 150),  # d7bd96
+    'Class6': (160, 126, 86),   # a07e56
+    'Class7': (130, 92, 67),    # 825c43
+    'Class8': (96, 65, 52),     # 604134
+    'Class9': (58, 49, 42),     # 3a312a
+    'Class10': (41, 36, 32),    # 292420
+}
+
+# Function to convert RGB tuple to hex color code
+def rgb_to_hex(rgb):
+    return '#{:02x}{:02x}{:02x}'.format(*rgb)
+
+# Mapping of Monk classes to colors using monk_scale
+monk_colors = {
+    '1': [rgb_to_hex(monk_scale['Class2']), rgb_to_hex(monk_scale['Class3']), rgb_to_hex(monk_scale['Class4'])],
+    '2': [rgb_to_hex(monk_scale['Class5']), rgb_to_hex(monk_scale['Class6'])],
+    '3': [rgb_to_hex(monk_scale['Class7']), rgb_to_hex(monk_scale['Class8'])],
+    '4': [rgb_to_hex(monk_scale['Class9']), rgb_to_hex(monk_scale['Class10'])],
+    'default': '#808080'  # Default color for unexpected classes
+}
+
+# Mapping of model's output classes to monk classes
+class_mapping = {
+    0: '1',  # Map model class 0 to monk class 1
+    1: '2',  # Map model class 1 to monk class 2
+    2: '3',  # Map model class 2 to monk class 3
+    3: '4',  # Map model class 3 to monk class 4
+    # Add more mappings if needed
+}
+
+# Function to load the model
+def load_model():
+    model_path = r"C:\Users\ramam\svm_model3.joblib"  # Adjust the path to your model
+    model = load(model_path)
+    return model
+
+# Function to parse face and extract skin region
+def parse_face(image):
+    # Ensure the image has 3 channels (RGB)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
+    image_data = np.array(image)
+    
+    # Check if the image has 3 channels
+    if image_data.shape[2] != 3:
+        raise ValueError("Image does not have 3 channels (RGB).")
+    
+    image_tensor = torch.from_numpy(image_data.astype('float32')).permute(2, 0, 1).unsqueeze(0).to(device)
+    faces = face_detector(image_tensor)
+
+    if faces:
+        parsed_faces = face_parser(image_tensor, faces)
+        if 'seg' in parsed_faces:
+            seg_logits = parsed_faces['seg']['logits']
+            seg_probs = torch.sigmoid(seg_logits)
+            binary_mask = seg_probs[0, 1, :, :] > 0.5
+            binary_mask = binary_mask.cpu().numpy()
+            binary_mask_3d = np.repeat(binary_mask[:, :, np.newaxis], 3, axis=2)
+            skin_region = image_data * binary_mask_3d
+            return skin_region.astype(np.uint8)
+    return None
+
+# Function to make predictions
+def classify_image(image, model):
+    parsed_image = parse_face(image)
+    if parsed_image is not None:
+        image_resized = resize(parsed_image, (128, 128), anti_aliasing=True)  # Resize to 128x128
+        image_reshaped = image_resized.reshape(1, -1)  # Reshape to match the model input
+        if image_reshaped.shape[1] == 49152:  # Check if resizing is correct
+            image_padded = np.pad(image_reshaped, ((0, 0), (0, 65536 - 49152)), 'constant')
+        else:
+            raise ValueError("Unexpected number of features after reshaping.")
+        prediction = model.predict(image_padded)
+        return prediction[0], parsed_image
+    else:
+        raise ValueError("Face parsing failed.")
+
+# Load the model
+model = load_model()
+
+# Function to display the Monk class color
+def display_monk_class_color(prediction):
+    st.write(f"Prediction: {prediction}")  # Debugging
+    monk_class = class_mapping.get(prediction, 'default')
+    colors = monk_colors.get(monk_class, monk_colors['default'])  # Default to gray if class not found
+    st.write(f"Monk Class: {monk_class}")
+    for color in colors:
+        st.markdown(f"<div style='width:100px; height:50px; background-color:{color};'></div>", unsafe_allow_html=True)
+
+# Streamlit app
+st.title('Skin Tone Classification')
+
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption='Uploaded Image.', use_column_width=True)
+
+    if st.button('Classify'):
+        try:
+            prediction, parsed_image = classify_image(image, model)
+            display_monk_class_color(prediction)
+            st.image(parsed_image, caption='Parsed Image.', use_column_width=True)
+        except ValueError as e:
+            st.error(f"Error: {e}")