Spaces:
No application file
No application file
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}") | |