Spaces:
Sleeping
Sleeping
import gradio as gr | |
import cv2 | |
import os | |
import mediapipe as mp | |
import numpy as np | |
import tempfile | |
from maps import * | |
from functions import * | |
class CosmeticInjectionVisualizer: | |
def __init__(self, muscles_map, tasks_map): | |
self.mp_face_mesh = mp.solutions.face_mesh | |
self.mp_drawing = mp.solutions.drawing_utils | |
self.mp_drawing_styles = mp.solutions.drawing_styles | |
self.muscles_map = muscles_map | |
self.tasks_map = tasks_map | |
def process_image(self, image, task_name): | |
frame_shape = image.shape | |
with self.mp_face_mesh.FaceMesh( | |
static_image_mode=True, | |
refine_landmarks=True, | |
max_num_faces=1, | |
min_detection_confidence=0.5) as face_mesh: | |
results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
if results.multi_face_landmarks: | |
for face_landmarks in results.multi_face_landmarks: | |
points = self.get_muscle_points(task_name, face_landmarks, frame_shape) | |
self.draw_muscle_points(image, points, face_landmarks, task_name) | |
return image | |
def draw_rounded_rectangle(self,image, start_point, end_point, color, thickness, radius): | |
top_left = start_point | |
bottom_right = end_point | |
if thickness < 0: # Filled rectangle | |
thickness = cv2.FILLED | |
# Draw filled rectangle with rounded corners | |
if thickness == cv2.FILLED: | |
# Top-left corner | |
cv2.ellipse(image, (top_left[0] + radius, top_left[1] + radius), (radius, radius), 180, 0, 90, color, -1) | |
# Top-right corner | |
cv2.ellipse(image, (bottom_right[0] - radius, top_left[1] + radius), (radius, radius), 270, 0, 90, color, -1) | |
# Bottom-left corner | |
cv2.ellipse(image, (top_left[0] + radius, bottom_right[1] - radius), (radius, radius), 90, 0, 90, color, -1) | |
# Bottom-right corner | |
cv2.ellipse(image, (bottom_right[0] - radius, bottom_right[1] - radius), (radius, radius), 0, 0, 90, color, -1) | |
# Top and bottom border | |
cv2.rectangle(image, (top_left[0] + radius, top_left[1]), (bottom_right[0] - radius, top_left[1] + radius), color, -1) | |
cv2.rectangle(image, (top_left[0] + radius, bottom_right[1] - radius), (bottom_right[0] - radius, bottom_right[1]), color, -1) | |
# Left and right border | |
cv2.rectangle(image, (top_left[0], top_left[1] + radius), (top_left[0] + radius, bottom_right[1] - radius), color, -1) | |
cv2.rectangle(image, (bottom_right[0] - radius, top_left[1] + radius), (bottom_right[0], bottom_right[1] - radius), color, -1) | |
# Center rectangle | |
cv2.rectangle(image, (top_left[0] + radius, top_left[1] + radius), (bottom_right[0] - radius, bottom_right[1] - radius), color, -1) | |
else: | |
# Top-left corner | |
cv2.ellipse(image, (top_left[0] + radius, top_left[1] + radius), (radius, radius), 180, 0, 90, color, thickness) | |
# Top-right corner | |
cv2.ellipse(image, (bottom_right[0] - radius, top_left[1] + radius), (radius, radius), 270, 0, 90, color, thickness) | |
# Bottom-left corner | |
cv2.ellipse(image, (top_left[0] + radius, bottom_right[1] - radius), (radius, radius), 90, 0, 90, color, thickness) | |
# Bottom-right corner | |
cv2.ellipse(image, (bottom_right[0] - radius, bottom_right[1] - radius), (radius, radius), 0, 0, 90, color, thickness) | |
# Top border | |
cv2.line(image, (top_left[0] + radius, top_left[1]), (bottom_right[0] - radius, top_left[1]), color, thickness) | |
# Bottom border | |
cv2.line(image, (top_left[0] + radius, bottom_right[1]), (bottom_right[0] - radius, bottom_right[1]), color, thickness) | |
# Left border | |
cv2.line(image, (top_left[0], top_left[1] + radius), (top_left[0], bottom_right[1] - radius), color, thickness) | |
# Right border | |
cv2.line(image, (bottom_right[0], top_left[1] + radius), (bottom_right[0], bottom_right[1] - radius), color, thickness) | |
def draw_muscle_points(self, image, points, face_landmarks, task, draw_background=False, verbose=False): | |
# Calculate bounding box of the face landmarks | |
x_min = min([landmark.x for landmark in face_landmarks.landmark]) * image.shape[1] | |
y_min = min([landmark.y for landmark in face_landmarks.landmark]) * image.shape[0] | |
x_max = max([landmark.x for landmark in face_landmarks.landmark]) * image.shape[1] | |
y_max = max([landmark.y for landmark in face_landmarks.landmark]) * image.shape[0] | |
face_width = x_max - x_min | |
face_height = y_max - y_min | |
# Determine text size and circle size based on face size relative to the image | |
scale_factor = 0.0005 * (face_width + face_height) | |
text_scale = scale_factor | |
thickness = int(2 * scale_factor) | |
margin = int(10 * scale_factor) | |
circle_radius = int(5 * scale_factor) | |
radius = int(10 * scale_factor) # For rounded corners | |
muscle_names = set() | |
for (x, y, muscle) in points: | |
# Draw the circle for the muscle point | |
cv2.circle(image, (x, y), circle_radius, (0, 0, 255), -1) | |
# Determine text size and background size | |
text = "3 U" | |
(text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, thickness) | |
text_x = x | |
text_y = y - 10 | |
text_x_end = text_x + text_width + 2 * margin | |
text_y_end = text_y - text_height - 2 * margin | |
if draw_background: | |
# Draw background rectangle with margins and rounded corners | |
self.draw_rounded_rectangle(image, (text_x - margin, text_y + margin), (text_x_end, text_y_end), (0, 0, 0), cv2.FILLED, radius) | |
if verbose: | |
# Draw the text on top of the rectangle | |
cv2.putText(image, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, text_scale, (0, 255, 0), thickness, cv2.LINE_AA) | |
muscle_names.add(muscle) | |
muscles = ','.join(muscle_names) | |
# Draw legend in the bottom-left corner with background | |
legend_texts = ['Total dose : 42','Name of patient: Julia Juila',f'Muscle : {muscles}', f'Task : {task}','Cosmetic App'] | |
legend_x = 10 | |
legend_y = image.shape[0] - 10 | |
legend_margin = 5 | |
legend_scale_factor = 1.5 * text_scale # Make legend text larger | |
legend_thickness = int(2 * legend_scale_factor) | |
legend_radius = int(10 * legend_scale_factor) | |
max_text_width = 0 | |
total_text_height = 0 | |
for text in legend_texts: | |
(text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, legend_thickness) | |
max_text_width = max(max_text_width, text_width) | |
total_text_height += text_height + 2 * legend_margin | |
legend_start_point = (legend_x - legend_margin, legend_y - total_text_height - legend_margin) | |
legend_end_point = (legend_x + max_text_width + legend_margin, legend_y + legend_margin) | |
if True: | |
self.draw_rounded_rectangle(image, legend_start_point, legend_end_point, (0, 0, 0), cv2.FILLED, legend_radius) | |
for text in legend_texts: | |
(text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, legend_thickness) | |
cv2.putText(image, text, (legend_x, legend_y), cv2.FONT_HERSHEY_SIMPLEX, legend_scale_factor, (255, 255, 255), legend_thickness, cv2.LINE_AA) | |
legend_y -= text_height + 2 * legend_margin | |
def get_muscle_points(self, task_name, face_landmarks, frame_shape): | |
if task_name not in self.tasks_map: | |
raise ValueError(f"Task '{task_name}' not found in tasks map.") | |
muscles_names = self.tasks_map[task_name]['muscles'] | |
points = [] | |
for muscle in muscles_names: | |
for region in self.muscles_map[muscle]: | |
if 'points' in self.muscles_map[muscle][region]: | |
for point_idx in self.muscles_map[muscle][region]['points']: | |
landmark = face_landmarks.landmark[point_idx] | |
x = int(landmark.x * frame_shape[1]) | |
y = int(landmark.y * frame_shape[0]) | |
points.append((x, y, muscle)) | |
else: | |
for subregion in self.muscles_map[muscle][region]: | |
for point_idx in self.muscles_map[muscle][region][subregion]['points']: | |
landmark = face_landmarks.landmark[point_idx] | |
x = int(landmark.x * frame_shape[1]) | |
y = int(landmark.y * frame_shape[0]) | |
points.append((x, y, muscle)) | |
return points | |
def process_video(self, video_path, task_name): | |
cap = cv2.VideoCapture(video_path) | |
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
output_path = temp_file.name | |
# Get the width and height of the frames | |
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
# Define the codec and create VideoWriter object | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Use appropriate codec for .mp4 files | |
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) | |
with self.mp_face_mesh.FaceMesh( | |
static_image_mode=False, | |
refine_landmarks=True, | |
max_num_faces=1, | |
min_detection_confidence=0.5, | |
min_tracking_confidence=0.5) as face_mesh: | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
frame_shape = frame.shape | |
results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) | |
if results.multi_face_landmarks: | |
for face_landmarks in results.multi_face_landmarks: | |
points = self.get_muscle_points(task_name, face_landmarks, frame_shape) | |
self.draw_muscle_points(frame, points, face_landmarks, task_name) | |
# Write the processed frame to the output video file | |
out.write(frame) | |
cap.release() | |
out.release() | |
cv2.destroyAllWindows() | |
return output_path | |
def process_webcam(self, task_name): | |
cap = cv2.VideoCapture(0) | |
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
output_path = temp_file.name | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
out = cv2.VideoWriter(output_path, fourcc, 20.0, (640, 480)) | |
with self.mp_face_mesh.FaceMesh( | |
static_image_mode=False, | |
refine_landmarks=True, | |
max_num_faces=1, | |
min_detection_confidence=0.5, | |
min_tracking_confidence=0.5) as face_mesh: | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
frame_shape = frame.shape | |
results = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) | |
if results.multi_face_landmarks: | |
for face_landmarks in results.multi_face_landmarks: | |
points = self.get_muscle_points(task_name, face_landmarks, frame_shape) | |
self.draw_muscle_points(frame, points, face_landmarks, task_name) | |
out.write(frame) | |
cap.release() | |
out.release() | |
cv2.destroyAllWindows() | |
return output_path | |
visualizer = CosmeticInjectionVisualizer(muscles_map, tasks_map) | |
def inference_image(image, task_name): | |
result_image = visualizer.process_image(image, task_name) | |
return result_image | |
def inference_video(video_path, task_name): | |
result_video_path = visualizer.process_video(video_path, task_name) | |
return result_video_path | |
def inference_webcam(task_name): | |
result_video_path = visualizer.process_webcam(task_name) | |
return result_video_path | |
task_names = list(tasks_map.keys()) | |
base_path=os.getcwd() | |
default_image_path = os.path.join(base_path,'image.jpg') | |
default_video_path = os.path.join(base_path,'video.mp4') | |
with gr.Blocks() as demo: | |
gr.Markdown("# Cosmetic Injection Visualizer") | |
with gr.Tabs(): | |
with gr.TabItem("Image"): | |
image_input = gr.Image(type="numpy", label="Input Image", value=default_image_path) | |
task_input_image = gr.Dropdown(choices=task_names, label="Task Name") | |
image_output = gr.Image(type="numpy", label="Output Image") | |
gr.Button("Process Image").click(inference_image, inputs=[image_input, task_input_image], outputs=image_output) | |
with gr.TabItem("Video"): | |
video_input = gr.Video(label="Input Video", value=default_video_path) | |
task_input_video = gr.Dropdown(choices=task_names, label="Task Name") | |
video_output = gr.Video(label="Output Video") | |
gr.Button("Process Video").click(inference_video, inputs=[video_input, task_input_video], outputs=video_output) | |
with gr.TabItem("Webcam"): | |
task_input_webcam = gr.Dropdown(choices=task_names, label="Task Name") | |
webcam_output = gr.Video(label="Output Video") | |
gr.Button("Process Webcam").click(inference_webcam, inputs=[task_input_webcam], outputs=webcam_output) | |
demo.launch() | |