Spaces:
Sleeping
Sleeping
import os | |
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' | |
os.environ['CUDA_VISIBLE_DEVICES'] = '0' | |
import tensorflow as tf | |
import tf_bodypix | |
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths | |
from tf_bodypix.draw import draw_poses | |
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 | |
import gradio as gr | |
import pandas as pd | |
# Load BodyPix model | |
bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16)) | |
rainbow = [ | |
[110, 64, 170], [143, 61, 178], [178, 60, 178], [210, 62, 167], | |
[238, 67, 149], [255, 78, 125], [255, 94, 99], [255, 115, 75], | |
[255, 140, 56], [239, 167, 47], [217, 194, 49], [194, 219, 64], | |
[175, 240, 91], [135, 245, 87], [96, 247, 96], [64, 243, 115], | |
[40, 234, 141], [28, 219, 169], [26, 199, 194], [33, 176, 213], | |
[47, 150, 224], [65, 125, 224], [84, 101, 214], [99, 81, 195] | |
] | |
def process_images(front_img, side_img, real_height_cm): | |
fimage_array = preprocessing.image.img_to_array(front_img) | |
simage_array = preprocessing.image.img_to_array(side_img) | |
# BodyPix prediction | |
frontresult = bodypix_model.predict_single(fimage_array) | |
sideresult = bodypix_model.predict_single(simage_array) | |
front_mask = frontresult.get_mask(threshold=0.75) | |
side_mask = sideresult.get_mask(threshold=0.75) | |
front_colored_mask = frontresult.get_colored_part_mask(front_mask, rainbow) | |
side_colored_mask = sideresult.get_colored_part_mask(side_mask, rainbow) | |
frontposes = frontresult.get_poses() | |
sideposes = sideresult.get_poses() | |
# Calculate body sizes | |
body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, sideposes, frontposes, real_height_cm, rainbow) | |
print(body_sizes) | |
measurements_df = pd.DataFrame([body_sizes]) if isinstance(body_sizes, dict) else pd.DataFrame(body_sizes) | |
# Save measurements to CSV | |
csv_file = "Body-measurement.csv" | |
if not os.path.exists(csv_file): | |
measurements_df.to_csv(csv_file, index=False) | |
else: | |
measurements_df.to_csv(csv_file, mode='a', header=False, index=False) | |
# Prepare measurements for display | |
measurement_display = measurements_df.to_html(index=False, justify="center", border=1) | |
print(measurement_display) | |
return f""" | |
<h3 style="text-align: center;">Body Measurements</h3> | |
<div style="text-align: center;"> | |
</div> | |
{measurement_display} | |
""" | |
# Create the Gradio interface | |
interface = gr.Interface( | |
fn=process_images, | |
inputs=[ | |
gr.Image(sources=["webcam", "upload"], type="numpy", label="Front Pose"), | |
gr.Image(sources=["webcam", "upload"], type="numpy", label="Side Pose"), | |
gr.Number(label="Enter Your Height (cm)") | |
], | |
outputs=gr.HTML(label="Measurement Results"), # Use HTML output to display the measurements | |
title="Body Sizing System Demo", | |
description="Capture two webcam images: Front View and Side View, and input your height in cm." | |
) | |
# Launch the app | |
interface.launch(share=True) |