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}")