SkinAiAnalyzer / main.py
AmirKaseb's picture
Update main.py
87b6396 verified
import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestRegressor
from skimage import filters
def calculate_redness_score(image_path):
img = cv2.imread(image_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(gray_img, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
if len(faces) == 0:
return 0, None
x, y, w, h = faces[0]
face_roi = img[y:y+h, x:x+w]
face_hsv = cv2.cvtColor(face_roi, cv2.COLOR_BGR2HSV)
lower_red = np.array([0, 70, 100])
upper_red = np.array([10, 255, 255])
mask = cv2.inRange(face_hsv, lower_red, upper_red)
redness_percentage = (np.count_nonzero(mask) / (w * h)) * 100
edge_mask = filters.sobel(mask)
segmented_image = cv2.watershed(face_roi, edge_mask)
regressor = RandomForestRegressor()
X = np.array([[x, y, w, h] for x, y, w, h in faces])
y = np.array([redness_percentage for _ in range(len(faces))])
regressor.fit(X, y)
plt.subplot(2, 2, 1)
plt.imshow(img_rgb)
plt.title('Original Image')
plt.subplot(2, 2, 2)
plt.imshow(cv2.cvtColor(cv2.bitwise_and(face_roi, face_roi, mask=mask), cv2.COLOR_BGR2RGB))
plt.title('Red Mask')
plt.subplot(2, 2, 3)
plt.imshow(edge_mask, cmap='gray')
plt.title('Edge Mask')
plt.subplot(2, 2, 4)
plt.imshow(segmented_image)
plt.title('Segmented Image')
plt.show()
return redness_percentage, mask
image_path = 'main-qimg-8f68ebd077d3e8c95f1af2a120adbe90-lq.jpg'
redness_score, red_mask = calculate_redness_score(image_path)
print(f'Redness Score: {redness_score:.2f}')
if red_mask is not None:
cv2.imwrite('red_mask.jpg', red_mask)