Spaces:
Running
Running
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}") |