diff --git a/config.yml b/config.yml index 6979d8e..91dd880 100644 --- a/config.yml +++ b/config.yml @@ -1,57 +1,15 @@ measurements: - - name: arm_length - landmarks: - - 11 - - 13 - - 15 - - - name: leg_length - landmarks: - - 23 - - 25 - - 27 - - name: shoulder_length - landmarks: - - 11 - - 12 - - - name: neck_to_hip_length - landmarks: - - 11 - - 23 - -#0 - nose -#1 - left eye (inner) -#2 - left eye -#3 - left eye (outer) -#4 - right eye (inner) -#5 - right eye -#6 - right eye (outer) -#7 - left ear -#8 - right ear -#9 - mouth (left) -#10 - mouth (right) -#11 - left shoulder -#12 - right shoulder -#13 - left elbow -#14 - right elbow -#15 - left wrist -#16 - right wrist -#17 - left pinky -#18 - right pinky -#19 - left index -#20 - right index -#21 - left thumb -#22 - right thumb -#23 - left hip -#24 - right hip -#25 - left knee -#26 - right knee -#27 - left ankle -#28 - right ankle -#29 - left heel -#30 - right heel -#31 - left foot index -#32 - right foot index - + landmarks: + - left_shoulder + - right_shoulder + - name: arm_length + landmarks: + - left_shoulder + - left_elbow + - left_wrist + - name: leg_length + landmarks: + - left_hip + - left_knee + - left_ankle diff --git a/dockerfile b/dockerfile deleted file mode 100644 index 9277476..0000000 --- a/dockerfile +++ /dev/null @@ -1,10 +0,0 @@ -FROM python:3.10-slim - -WORKDIR /app - -COPY . /app - -RUN pip install --no-cache-dir -r requirements.txt -RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y - -ENTRYPOINT ["sh", "-c", "python3 landmarks.py \"$@\" 2>&1 | grep -vE '^(WARNING:|I0000|INFO:)' | awk '{print}'", "--"] diff --git a/landmarks.js b/landmarks.js new file mode 100644 index 0000000..cbbb3e9 --- /dev/null +++ b/landmarks.js @@ -0,0 +1,218 @@ +const fs = require('fs'); +const path = require('path'); +const ArgumentParser = require('argparse'); +const cv = require('opencv'); +const yaml = require('js-yaml'); +const { Pose, POSE_LANDMARKS } = require('@mediapipe/pose'); + +const logging = console; +const warnings = console; + +class Landmarker { + static resizedHeight = 256; + static resizedWidth = 256; + + constructor() { + this.args = this.parseArgs(); + this.measurements = this.loadLandmarks(); + if (!this.args.frontImage) { + throw new Error("Front image needs to be passed"); + } + if (!this.args.sideImage) { + throw new Error("Side image needs to be passed"); + } + + this.frontImage = cv.imread(this.args.frontImage); + this.sideImage = cv.imread(this.args.sideImage); + + this.frontImageResized = this.frontImage.resize(Landmarker.resizedHeight, Landmarker.resizedWidth); + this.sideImageResized = this.sideImage.resize(Landmarker.resizedHeight, Landmarker.resizedWidth); + + this.distances = {}; + + this.personHeight = this.args.personHeight; + this.pixelHeight = this.args.pixelHeight; + + this.pose = new Pose({ + locateFile: (file) => { + return `https://cdn.jsdelivr.net/npm/@mediapipe/pose/${file}`; + } + }); + + this.landmarksIndices = [ + POSE_LANDMARKS.LEFT_SHOULDER, + POSE_LANDMARKS.RIGHT_SHOULDER, + POSE_LANDMARKS.LEFT_ELBOW, + POSE_LANDMARKS.RIGHT_ELBOW, + POSE_LANDMARKS.LEFT_WRIST, + POSE_LANDMARKS.RIGHT_WRIST, + POSE_LANDMARKS.LEFT_HIP, + POSE_LANDMARKS.RIGHT_HIP, + POSE_LANDMARKS.LEFT_KNEE, + POSE_LANDMARKS.RIGHT_KNEE, + POSE_LANDMARKS.LEFT_ANKLE, + POSE_LANDMARKS.RIGHT_ANKLE + ]; + } + + loadLandmarks() { + const file = fs.readFileSync(this.args.yamlFile, 'utf8'); + const landmarksData = yaml.load(file); + const measurements = {}; + for (const measurement of landmarksData.measurements) { + measurements[measurement.name] = measurement.landmarks; + } + return measurements; + } + + parseArgs() { + const parser = new ArgumentParser(); + parser.add_argument('--front', { dest: 'frontImage', type: 'str', help: 'Front image' }); + parser.add_argument('--side', { dest: 'sideImage', type: 'str', help: 'Side image' }); + parser.add_argument('--poseDetectionConfidence', { dest: 'poseDetectionConfidence', default: 0.5, type: 'float', help: 'Confidence score for pose detection' }); + parser.add_argument('--poseTrackingConfidence', { dest: 'poseTrackingConfidence', default: 0.5, type: 'float', help: 'Confidence score for pose tracking' }); + parser.add_argument('--personHeight', { dest: 'personHeight', type: 'int', help: 'Person height in cm' }); + parser.add_argument('--pixelHeight', { dest: 'pixelHeight', type: 'int', help: 'Pixel height of person' }); + parser.add_argument('--measurement', { dest: 'measurement', nargs: '+', type: 'str', help: 'Type of measurement' }); + parser.add_argument('--yamlFile', { dest: 'yamlFile', type: 'str', help: 'Path to the YAML file containing landmarks' }); + return parser.parse_args(); + } + + async run() { + await this.pose.initialize(); + const { frontResults, sideResults } = await this.processImages(); + + this.getCenterTopPoint(sideResults); + + const table = []; + if (this.args.measurement) { + for (const m of this.args.measurement) { + if (!this.measurements[m]) { + throw new Error("Incorrect input (input not present in config.yml)"); + } else { + const distance = this.calculateDistanceBetweenLandmarks(frontResults, m); + table.push([m, distance]); + } + } + } else { + for (const m in this.measurements) { + const distance = this.calculateDistanceBetweenLandmarks(frontResults, m); + table.push([m, distance]); + } + } + + console.table(table); + + this.pose.close(); + } + + async processImages() { + const frontResults = await this.pose.estimatePoses(this.frontImageResized); + const sideResults = await this.pose.estimatePoses(this.sideImageResized); + + this.sideImageKeypoints = this.sideImageResized.copy(); + this.frontImageKeypoints = this.frontImageResized.copy(); + + if (frontResults[0].landmarks) { + this.drawLandmarks(this.frontImageKeypoints, frontResults[0].landmarks, this.landmarksIndices); + } + if (sideResults[0].landmarks) { + this.drawLandmarks(this.sideImageKeypoints, sideResults[0].landmarks, this.landmarksIndices); + } + return { + frontResults: frontResults[0], + sideResults: sideResults[0] + }; + } + + pixelToMetricRatio() { + this.pixelHeight = this.pixelDistance * 2; + const pixelToMetricRatio = this.personHeight / this.pixelHeight; + logging.debug("pixelToMetricRatio %s", pixelToMetricRatio); + return pixelToMetricRatio; + } + + drawLandmarks(image, landmarks, indices) { + for (const idx of indices) { + const landmark = landmarks[idx]; + const h = image.rows; + const w = image.cols; + const cx = Math.round(landmark.x * w); + const cy = Math.round(landmark.y * h); + this.circle(image, cx, cy); + } + } + + circle(image, cx, cy) { + return image.drawCircle(new cv.Point(cx, cy), 2, new cv.Vec(255, 0, 0), -1); + } + + calculateDistanceBetweenLandmarks(frontResults, measurementName) { + if (!frontResults.landmarks) { + return; + } + + const landmarks = frontResults.landmarks; + const landmarkNames = this.measurements[measurementName]; + + let totalDistance = 0; + for (let i = 0; i < landmarkNames.length - 1; i++) { + const current = landmarks[landmarkNames[i]]; + const next = landmarks[landmarkNames[i + 1]]; + const pixelDistance = this.euclideanDistance( + current.x * Landmarker.resizedWidth, + current.y * Landmarker.resizedHeight, + next.x * Landmarker.resizedWidth, + next.y * Landmarker.resizedHeight + ); + const realDistance = pixelDistance * this.pixelToMetricRatio(); + totalDistance += realDistance; + } + return totalDistance; + } + + euclideanDistance(x1, y1, x2, y2) { + return Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2); + } + + getCenterTopPoint(sideResults) { + const grayImage = this.sideImageKeypoints.cvtColor(cv.COLOR_BGR2GRAY); + const blurredImage = grayImage.gaussianBlur(new cv.Size(5, 5), 0); + const roi = blurredImage.getRegion(new cv.Rect(0, 0, this.sideImageResized.cols, Math.floor(this.sideImageResized.rows / 2))); + this.edges = roi.canny(50, 150); + const contours = this.edges.findContours(cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE); + + let xt, yt; + this.topmostPoint = null; + + if (contours.length > 0) { + const largestContour = contours.sort((a, b) => b.area - a.area)[0]; + this.topmostPoint = largestContour.minBy(pt => pt.y); + xt = this.topmostPoint.x; + yt = this.topmostPoint.y; + + this.circle(this.sideImageKeypoints, xt, yt); + } + + let xc, yc; + if (sideResults.landmarks) { + const leftHip = sideResults.landmarks[POSE_LANDMARKS.LEFT_HIP]; + const rightHip = sideResults.landmarks[POSE_LANDMARKS.RIGHT_HIP]; + const centerPoint = [ + (leftHip.x + rightHip.x) / 2, + (leftHip.y + rightHip.y) / 2 + ]; + xc = Math.round(centerPoint[0] * this.sideImageResized.cols); + yc = Math.round(centerPoint[1] * this.sideImageResized.rows); + this.circle(this.sideImageKeypoints, xc, yc); + } + + this.pixelDistance = this.euclideanDistance(xt, yt, xc, yc); + this.pixelDistance *= 2; + } +} + +(async () => { + const landmarker = new Landmarker(); + await landmarker.run(); +})(); diff --git a/landmarks.py b/landmarks.py index a2d133a..8f6280b 100644 --- a/landmarks.py +++ b/landmarks.py @@ -18,6 +18,41 @@ warnings.filterwarnings( module="google.protobuf", ) +LANDMARK_NAME_TO_INDEX = { + "nose": 0, + "left_eye_inner": 1, + "left_eye": 2, + "left_eye_outer": 3, + "right_eye_inner": 4, + "right_eye": 5, + "right_eye_outer": 6, + "left_ear": 7, + "right_ear": 8, + "mouth_left": 9, + "mouth_right": 10, + "left_shoulder": 11, + "right_shoulder": 12, + "left_elbow": 13, + "right_elbow": 14, + "left_wrist": 15, + "right_wrist": 16, + "left_pinky": 17, + "right_pinky": 18, + "left_index": 19, + "right_index": 20, + "left_thumb": 21, + "right_thumb": 22, + "left_hip": 23, + "right_hip": 24, + "left_knee": 25, + "right_knee": 26, + "left_ankle": 27, + "right_ankle": 28, + "left_heel": 29, + "right_heel": 30, + "left_foot_index": 31, + "right_foot_index": 32, +} class Landmarker: @@ -50,18 +85,18 @@ class Landmarker: ) self.landmarks_indices = [ - pose.PoseLandmark.LEFT_SHOULDER.value, - pose.PoseLandmark.RIGHT_SHOULDER.value, - pose.PoseLandmark.LEFT_ELBOW.value, - pose.PoseLandmark.RIGHT_ELBOW.value, - pose.PoseLandmark.LEFT_WRIST.value, - pose.PoseLandmark.RIGHT_WRIST.value, - pose.PoseLandmark.LEFT_HIP.value, - pose.PoseLandmark.RIGHT_HIP.value, - pose.PoseLandmark.LEFT_KNEE.value, - pose.PoseLandmark.RIGHT_KNEE.value, - pose.PoseLandmark.LEFT_ANKLE.value, - pose.PoseLandmark.RIGHT_ANKLE.value, + LANDMARK_NAME_TO_INDEX["left_shoulder"], + LANDMARK_NAME_TO_INDEX["right_shoulder"], + LANDMARK_NAME_TO_INDEX["left_elbow"], + LANDMARK_NAME_TO_INDEX["right_elbow"], + LANDMARK_NAME_TO_INDEX["left_wrist"], + LANDMARK_NAME_TO_INDEX["right_wrist"], + LANDMARK_NAME_TO_INDEX["left_hip"], + LANDMARK_NAME_TO_INDEX["right_hip"], + LANDMARK_NAME_TO_INDEX["left_knee"], + LANDMARK_NAME_TO_INDEX["right_knee"], + LANDMARK_NAME_TO_INDEX["left_ankle"], + LANDMARK_NAME_TO_INDEX["right_ankle"], ] def load_landmarks(self): @@ -69,7 +104,7 @@ class Landmarker: landmarks_data = yaml.safe_load(file) measurements = {} for measurement in landmarks_data["measurements"]: - measurements[measurement["name"]] = measurement["landmarks"] + measurements[measurement["name"]] = [LANDMARK_NAME_TO_INDEX[l] for l in measurement["landmarks"]] return measurements def parse_args(self): @@ -257,66 +292,41 @@ class Landmarker: ] self.edges = cv2.Canny(roi, 50, 150) contours, _ = cv2.findContours( - self.edges, - cv2.RETR_EXTERNAL, + self.edges.copy(), + cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE, ) - xt, yt = None, None - self.topmost_point = None + max_contour = max(contours, key=cv2.contourArea) + rect = cv2.minAreaRect(max_contour) + box = cv2.boxPoints(rect) + box = sorted( + list(box), + key=lambda p: p[1], + ) + top_point = min( + box[0], + box[1], + key=lambda p: p[0], + ) - if contours: - largest_contour = max( - contours, - key=cv2.contourArea, - ) - self.topmost_point = tuple(largest_contour[largest_contour[:, :, 1].argmin()][0]) - xt, yt = self.topmost_point - - cv2.circle( - self.side_image_keypoints, - (xt, yt), - 2, - (255, 255, 0), - -1, - ) - xc, yc = None, None - landmarks = side_results.pose_landmarks.landmark - - if side_results.pose_landmarks: - left_hip = landmarks[pose.PoseLandmark.LEFT_HIP.value] - right_hip = landmarks[pose.PoseLandmark.RIGHT_HIP.value] - center_point = ( - (left_hip.x + right_hip.x) / 2, - (left_hip.y + right_hip.y) / 2, - ) - center_point = ( - int(center_point[0] * self.side_image_resized.shape[1]), - int(center_point[1] * self.side_image_resized.shape[0]), - ) - xc, yc = center_point - self.circle( - self.side_image_keypoints, - xc, - yc, - ) - - self.pixel_distance = self.euclidean_distance(xc, yc, xt, yt) - logging.debug( - "top_center_pixel_distance: %s", - self.pixel_distance, - ) - self.pixel_height = self.pixel_distance * 2 - logging.debug( - "pixel height: %s ", - self.pixel_height, - ) - self.distance = self.euclidean_distance(xc, yc, xt, yt) * self.pixel_to_metric_ratio() - return self.distance - - -l = Landmarker() -try: - l.run() -except Exception as e: - print(e) + left_hip = side_results.pose_landmarks.landmark[LANDMARK_NAME_TO_INDEX["left_hip"]] + right_hip = side_results.pose_landmarks.landmark[LANDMARK_NAME_TO_INDEX["right_hip"]] + + center_x = (left_hip.x + right_hip.x) / 2 + center_y = (left_hip.y + right_hip.y) / 2 + + center_x, center_y = ( + int(center_x * self.resized_width), + int(center_y * self.resized_height), + ) + + self.pixel_distance = self.euclidean_distance( + top_point[0], + top_point[1], + center_x, + center_y, + ) +if __name__ == "__main__": + landmarker = Landmarker() + landmarker.run()