SalehAhmad's picture
all
318f7bc
raw
history blame
1.86 kB
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import torch
torch.set_default_device('cuda:0')
import tensorflow as tf
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
from tf_bodypix.draw import draw_poses # utility function using OpenCV
from tensorflow.keras import preprocessing
import cv2
import json
from matplotlib import pyplot as plt
import numpy as np
from calculations import measure_body_sizes
bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))
input_path = 'input1/files/6'
testfile = 'front_img.jpg'
output_path = 'output'
image = preprocessing.image.load_img(input_path+'/'+testfile)
# image converted to image array
image_array = preprocessing.image.img_to_array(image)
# bodypix prediction
result = bodypix_model.predict_single(image_array)
mask = result.get_mask(threshold=0.75)
preprocessing.image.save_img(f'{output_path}/bodypix-mask.jpg',mask)
colored_mask = result.get_colored_part_mask(mask)
print(np.array(image).shape)
print(colored_mask.shape)
preprocessing.image.save_img(f'{output_path}/bodypix-colored-mask.jpg',colored_mask)
poses = result.get_poses()
image_with_poses = draw_poses(
image_array.copy(), # create a copy to ensure we are not modifing the source image
poses,
keypoints_color=(255, 100, 100),
skeleton_color=(100, 100, 255)
)
# print(poses)
preprocessing.image.save_img(f'{output_path}/bodypix-poses.jpg', image_with_poses)
real_height_cm = 155.0 # Replace with the real height in cm
body_sizes = 3 #measure_body_sizes(poses, real_height_cm)
print(body_sizes)
# Define the file name
file_name = './output/measurements.json'
# Open the file in write mode and save the dictionary as JSON
with open(file_name, 'w') as json_file:
json.dump(body_sizes, json_file, indent=4)
print(f"body_sizes saved to {file_name}")