This commit is contained in:
aparnah 2024-08-09 15:32:34 +05:30
parent 5856721002
commit 1b0fb383c6
2 changed files with 265 additions and 192 deletions

View File

@ -1,218 +1,289 @@
const fs = require('fs'); const fs = require("fs");
const path = require('path'); const path = require("path");
const ArgumentParser = require('argparse'); const { ArgumentParser } = require("argparse");
const cv = require('opencv'); const cv = require("@techstark/opencv-js");
const yaml = require('js-yaml'); const yaml = require("js-yaml");
const { Pose, POSE_LANDMARKS } = require('@mediapipe/pose'); const { Pose, POSE_LANDMARKS } = require("@mediapipe/pose");
const logging = console; const logging = console;
const warnings = console; const warnings = console;
class Landmarker { class Landmarker {
static resizedHeight = 256; static resizedHeight = 256;
static resizedWidth = 256; static resizedWidth = 256;
constructor() { constructor() {
this.args = this.parseArgs(); this.args = this.parseArgs();
this.measurements = this.loadLandmarks(); this.measurements = this.loadLandmarks();
if (!this.args.frontImage) { if (!this.args.frontImage) {
throw new Error("Front image needs to be passed"); throw new Error("Front image needs to be passed");
} }
if (!this.args.sideImage) { if (!this.args.sideImage) {
throw new Error("Side image needs to be passed"); 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() { this.frontImage = cv.imread(this.args.frontImage);
const file = fs.readFileSync(this.args.yamlFile, 'utf8'); this.sideImage = cv.imread(this.args.sideImage);
const landmarksData = yaml.load(file);
const measurements = {}; this.frontImageResized = cv.resize(
for (const measurement of landmarksData.measurements) { this.frontImage,
measurements[measurement.name] = measurement.landmarks; new cv.Size(Landmarker.resizedWidth, Landmarker.resizedHeight),
} );
return measurements; this.sideImageResized = cv.resize(
this.sideImage,
new cv.Size(Landmarker.resizedWidth, Landmarker.resizedHeight),
);
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() { parseArgs() {
const parser = new ArgumentParser(); const parser = new ArgumentParser({
parser.add_argument('--front', { dest: 'frontImage', type: 'str', help: 'Front image' }); description: "Process images and calculate measurements",
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("--front", {
parser.add_argument('--poseTrackingConfidence', { dest: 'poseTrackingConfidence', default: 0.5, type: 'float', help: 'Confidence score for pose tracking' }); dest: "frontImage",
parser.add_argument('--personHeight', { dest: 'personHeight', type: 'int', help: 'Person height in cm' }); required: true,
parser.add_argument('--pixelHeight', { dest: 'pixelHeight', type: 'int', help: 'Pixel height of person' }); help: "Path to the front image",
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' }); parser.add_argument("--side", {
return parser.parse_args(); dest: "sideImage",
} required: true,
help: "Path to the 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",
required: true,
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",
required: true,
help: "Path to the YAML file containing landmarks",
});
return parser.parse_args();
}
async run() { async run() {
await this.pose.initialize(); await this.pose.initialize();
const { frontResults, sideResults } = await this.processImages(); const { frontResults, sideResults } = await this.processImages();
this.getCenterTopPoint(sideResults); this.getCenterTopPoint(sideResults);
const table = []; const table = [];
if (this.args.measurement) { if (this.args.measurement) {
for (const m of this.args.measurement) { for (const m of this.args.measurement) {
if (!this.measurements[m]) { if (!this.measurements[m]) {
throw new Error("Incorrect input (input not present in config.yml)"); throw new Error("Incorrect input (input not present in config.yml)");
} else {
const distance = this.calculateDistanceBetweenLandmarks(frontResults, m);
table.push([m, distance]);
}
}
} else { } else {
for (const m in this.measurements) { const distance = this.calculateDistanceBetweenLandmarks(
const distance = this.calculateDistanceBetweenLandmarks(frontResults, m); frontResults,
table.push([m, distance]); m,
} );
table.push([m, distance]);
} }
}
console.table(table); } else {
for (const m in this.measurements) {
this.pose.close(); const distance = this.calculateDistanceBetweenLandmarks(
frontResults,
m,
);
table.push([m, distance]);
}
} }
async processImages() { console.table(table);
const frontResults = await this.pose.estimatePoses(this.frontImageResized);
const sideResults = await this.pose.estimatePoses(this.sideImageResized);
this.sideImageKeypoints = this.sideImageResized.copy(); this.pose.close();
this.frontImageKeypoints = this.frontImageResized.copy(); }
if (frontResults[0].landmarks) { async processImages() {
this.drawLandmarks(this.frontImageKeypoints, frontResults[0].landmarks, this.landmarksIndices); const frontResults = await this.pose.estimatePoses(this.frontImageResized);
} const sideResults = await this.pose.estimatePoses(this.sideImageResized);
if (sideResults[0].landmarks) {
this.drawLandmarks(this.sideImageKeypoints, sideResults[0].landmarks, this.landmarksIndices); this.sideImageKeypoints = this.sideImageResized.clone();
} this.frontImageKeypoints = this.frontImageResized.clone();
return {
frontResults: frontResults[0], if (frontResults[0].landmarks) {
sideResults: sideResults[0] 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() {
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) {
cv.circle(image, new cv.Point(cx, cy), 2, new cv.Scalar(255, 0, 0), -1);
}
calculateDistanceBetweenLandmarks(frontResults, measurementName) {
if (!frontResults.landmarks) {
return;
} }
pixelToMetricRatio() { const landmarks = frontResults.landmarks;
this.pixelHeight = this.pixelDistance * 2; const landmarkNames = this.measurements[measurementName];
const pixelToMetricRatio = this.personHeight / this.pixelHeight;
logging.debug("pixelToMetricRatio %s", pixelToMetricRatio); let totalDistance = 0;
return pixelToMetricRatio; 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 = cv.cvtColor(this.sideImageKeypoints, cv.COLOR_BGR2GRAY);
const blurredImage = cv.GaussianBlur(grayImage, new cv.Size(5, 5), 0);
const roi = blurredImage.roi(
new cv.Rect(
0,
0,
this.sideImageResized.cols,
Math.floor(this.sideImageResized.rows / 2),
),
);
this.edges = cv.Canny(roi, 50, 150);
const contours = this.edges.findContours(
cv.RETR_EXTERNAL,
cv.CHAIN_APPROX_SIMPLE,
);
let xt, yt;
this.topmostPoint = null;
for (const contour of contours) {
const [xt, yt] = contour.minEnclosingCircle();
if (this.topmostPoint === null || yt < this.topmostPoint[1]) {
this.topmostPoint = [xt, yt];
}
} }
drawLandmarks(image, landmarks, indices) { const { x, y } = sideResults.landmarks[POSE_LANDMARKS.NOSE];
for (const idx of indices) { const centerPoint = [
const landmark = landmarks[idx]; x * Landmarker.resizedWidth,
const h = image.rows; y * Landmarker.resizedHeight,
const w = image.cols; ];
const cx = Math.round(landmark.x * w); this.pixelHeight = Math.abs(centerPoint[1] - this.topmostPoint[1]);
const cy = Math.round(landmark.y * h);
this.circle(image, cx, cy);
}
}
circle(image, cx, cy) { cv.circle(
return image.drawCircle(new cv.Point(cx, cy), 2, new cv.Vec(255, 0, 0), -1); this.sideImageKeypoints,
} new cv.Point(centerPoint[0], centerPoint[1]),
2,
calculateDistanceBetweenLandmarks(frontResults, measurementName) { new cv.Scalar(255, 0, 0),
if (!frontResults.landmarks) { -1,
return; );
} cv.circle(
this.sideImageKeypoints,
const landmarks = frontResults.landmarks; new cv.Point(this.topmostPoint[0], this.topmostPoint[1]),
const landmarkNames = this.measurements[measurementName]; 2,
new cv.Scalar(255, 0, 0),
let totalDistance = 0; -1,
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();
const landmarker = new Landmarker(); landmarker.run().catch((error) => {
await landmarker.run(); console.error(error);
})(); });

View File

@ -189,6 +189,7 @@ class Landmarker:
tablefmt="plain", tablefmt="plain",
) )
print(output) print(output)
self.pose.close() self.pose.close()
def process_images(self): def process_images(self):
@ -279,6 +280,7 @@ class Landmarker:
distance = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) distance = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
return distance return distance
def get_center_top_point(self, side_results): def get_center_top_point(self, side_results):
gray_image = cv2.cvtColor( gray_image = cv2.cvtColor(
self.side_image_keypoints, self.side_image_keypoints,