Upload 5 files

app.py

@@ -0,0 +1,282 @@
+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()

functions.py

@@ -0,0 +1,57 @@
+import cv2
+import numpy as np
+import os 
+
+def is_bgr(image):
+    """
+    Heuristic check to see if an image is in BGR format.
+    This function checks if the values at a specific pixel match the expected BGR to RGB order.
+    """
+    # Choose a pixel to check (e.g., the top-left corner)
+    pixel = image[0, 0]
+    
+    # Check if the pixel values match BGR order
+    # This is a heuristic; it assumes the color channels have distinct values.
+    return pixel[0] > pixel[1] and pixel[1] > pixel[2]
+
+def convert_bgr_to_rgb(image):
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+def convert_rgb_to_bgr(image):
+    return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+
+def crop_video(input_path, output_path, x, y, width, height):
+    # Open the input video
+    cap = cv2.VideoCapture(input_path)
+
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return
+
+    # Get the width and height of the frames
+    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    # Define the codec and create VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can also use other codecs
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Crop the frame
+        cropped_frame = frame[y:y+height, x:x+width]
+
+        # Write the cropped frame to the output video file
+        out.write(cropped_frame)
+
+    # Release everything if job is finished
+    cap.release()
+    out.release()
+    cv2.destroyAllWindows()
+
+    print(f"Cropping completed. Saved to {output_path}")

maps.py

@@ -0,0 +1,241 @@
+muscles_map = {
+    'Corrugator Supercilii':{
+        'right':{
+            'points':[70,63,105,66,107],
+            'description':'',
+        },
+        'left':{
+            'points':[336,296,334,293,300],
+            'description':'',
+        },
+    },
+    'Procerus':{
+        'center':{
+            'points':[9,8],
+            'description':'',
+        },
+    },
+    'Medial Orbicularis Oculi':{
+        'right':{
+            'points':[221,189],
+            'description':'',
+        },
+        'left':{
+            'points':[413,441],
+            'description':'',
+        },
+    },
+    'Frontalis':{
+        'right':{
+            'superior':{
+                 'points':[104,69,108],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[103,67,109],
+                  'description':'',
+            },
+
+        },
+        'center':{
+            'superior':{
+                 'points':[10],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[151],
+                  'description':'',
+            },
+
+        },
+        'left':{
+            'superior':{
+                 'points':[338,297,332],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[337,299,333],
+                  'description':'',
+            },
+        },
+    },
+    'Lateral Orbicularis Oculi':{
+        'right':{
+            'internal':{
+                'points':[156,143,116],
+                'description':'',
+            },
+            'external':{
+                'points':[139,34,227],
+                'description':'',
+            },
+
+        },
+        'left':{
+            'internal':{
+                'points':[383,372,345],
+                'description':"Crow's feet",
+            },
+            'external':{
+                'points':[368,264,447],
+                'description':"Crow's feet",
+            },
+        },
+    },
+    'Lateral Superior Orbicularis Oculi':{
+        'right':{
+            'points':[224,225,113],
+            'description':"Eyebrow lift",
+        },
+        'left':{
+            'points':[444,445,342],
+            'description':"Eyebrow lift",
+        },
+    },
+    'Nasalis':{
+        'right':{
+            'points':[196,174,217],
+            'description':"Rabbit's lines",
+        },
+        'left':{
+            'points':[419,399,437],
+            'description':"Rabbit's lines",
+        },
+    },
+    'Dilator nasalis':{
+        'right':{
+            'points':[134,131],
+            'description':"Reduce nostril size",
+        },
+        'left':{
+            'points':[363,360],
+            'description':"Reduce nostril size",
+        },
+    },
+    'Depressor septi nasi':{
+        'center':{
+            'points':[2],
+            'description':"Nose lift",
+        },
+    },
+    'Depressor anguli oris':{
+        'right':{
+            'points':[202,169,149],
+            'description':"corner lip lift",
+        },
+        'left':{
+            'points':[422,394,378],
+            'description':"corner lip lift",
+        },
+    },
+    'Orbicularis':{
+        'right':{
+            'superior':{
+                 'points':[61,185,40,39,37],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[61,146,91,181,84],
+                  'description':'',
+            },
+
+        },
+        'center':{
+            'superior':{
+                 'points':[0],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[17],
+                  'description':'',
+            },
+
+        },
+        'left':{
+            'superior':{
+                 'points':[267,269,270,409,291],
+                  'description':'',
+            },
+            'inferior':{
+                 'points':[314,405,321,375,291],
+                  'description':'',
+            },
+        },
+    },
+    'Levator labii superioris alaeque nasi':{
+        'right':{
+            'points':[203],
+            'description':"gummy smile",
+        },
+        'left':{
+            'points':[423],
+            'description':"gummy smile",
+        },
+    },
+    'Mentalis':{
+        'right':{
+            'points':[171,148],
+            'description':"Chin crease",
+        },
+        'center':{
+            'points':[152],
+            'description':"Chin crease",
+        },
+        'left':{
+            'points':[396,377],
+            'description':"Chin crease",
+        },
+    },
+}
+
+
+tasks_map = {
+    "Gallabela crease":{
+        'muscles':['Corrugator Supercilii','Procerus'],
+        'description':"",
+    },
+    "Forehead crease":{
+        'muscles':['Frontalis'],
+        'description':"",
+    },
+    "Crow's feet crease":{
+        'muscles':['Lateral Orbicularis Oculi'],
+        'description':"",
+    },
+    "Lift of eyebrows":{
+        'muscles':['Lateral Superior Orbicularis Oculi'],
+        'description':"",
+    },
+    "Chemical lift of eyebrows":{
+        'muscles':['Lateral Superior Orbicularis Oculi','Corrugator Supercilii','Procerus'],
+        'description':"",
+    },
+    "Nose crease":{
+        'muscles':['Nasalis'],
+        'description':"",
+    },
+    "Lift of nose":{
+        'muscles':['Depressor septi nasi'],
+        'description':"",
+    },
+    "Reducing the size of nostrils":{
+        'muscles':['Dilator nasalis'],
+        'description':"",
+    },
+    "Lift of the corner of the lips":{
+        'muscles':['Depressor anguli oris'],
+        'description':"",
+    },
+    "Ciggarete crease":{
+        'muscles':['Orbicularis'],
+        'description':"",
+    },
+    "Gummy smile":{
+        'muscles':['Levator labii superioris alaeque nasi'],
+        'description':"",
+    },
+    "Chin crease":{
+        'muscles':['Mentalis'],
+        'description':"",
+    },
+}

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+opencv-python-headless
+mediapipe
+numpy