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() | |