Cosmetic_inj / app.py
Medvira's picture
Upload 5 files
d936db7 verified
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()