Hemant0000's picture
Update app.py
d41a7bc verified
import cv2
import numpy as np
import matplotlib.pyplot as plt
import mediapipe as mp
import gradio as gr
import mediapipe as mp
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
# classification head to MobileNetV2
base_model = MobileNetV2(weights="imagenet", include_top=False, input_shape=(128, 128, 3))
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1, activation='sigmoid')(x) # Sigmoid activation for binary classification (0 or 1)
model = Model(inputs=base_model.input, outputs=x)
def classify_gender(image):
# Resize image to 128x128 for MobileNetV2
img_resized = cv2.resize(image, (128, 128))
img_array = img_to_array(img_resized)
img_array = preprocess_input(img_array)
img_array = np.expand_dims(img_array, axis=0)
# Predict gender (assuming 0: Male, 1: Female)
prediction = model.predict(img_array)
gender = "Male" if prediction[0][0] < 0.5 else "Female"
return gender
# Initialize Mediapipe Pose model
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True)
# Function to estimate height using dynamic calibration based on shoulder width
def estimate_real_world_height_and_area(image):
height, width, _ = image.shape
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
if not results.pose_landmarks:
print("No landmarks detected.")
return None, None, None # Return None if no landmarks detected
landmarks = results.pose_landmarks.landmark
# Reference average shoulder width (in cm)
AVERAGE_SHOULDER_WIDTH_CM = 45
# Points for height (nose to ankle) and shoulder width
nose = landmarks[mp_pose.PoseLandmark.NOSE]
left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER]
right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER]
left_ankle = landmarks[mp_pose.PoseLandmark.LEFT_ANKLE]
right_ankle = landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE]
# Calculate pixel height between nose (head) and ankle
head_y = int(nose.y * height)
ankle_y = int((left_ankle.y + right_ankle.y) / 2 * height)
pixel_height = abs(ankle_y - head_y)
# Calculate shoulder width in pixels
shoulder_width_pixels = abs(left_shoulder.x - right_shoulder.x) * width
# Determine the pixel-to-cm ratio using shoulder width
pixel_to_cm_ratio = AVERAGE_SHOULDER_WIDTH_CM / shoulder_width_pixels if shoulder_width_pixels != 0 else 0
# Calculate real-world height using this ratio
real_world_height_cm = pixel_height * pixel_to_cm_ratio
# Estimate torso area in pixels
torso_width = abs(left_shoulder.x - right_shoulder.x) * width
torso_height = abs(left_shoulder.y - left_ankle.y) * height
torso_area_pixels = torso_width * torso_height
real_torso_area = torso_area_pixels * (pixel_to_cm_ratio ** 2)
return real_world_height_cm, real_torso_area, landmarks
# Main Gradio prediction function with scaling and plotting
def predict(image):
try:
# Classify Gender
gender = classify_gender(image) # Assuming this function is defined
# Estimate height and area
real_height, real_torso_area, landmarks = estimate_real_world_height_and_area(image)
if real_height is None or real_torso_area is None:
return "No landmarks detected; please use a clearer image.", None, None, None
# Plot with scale
fig, ax = plt.subplots(figsize=(10, 10))
ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# Plot landmarks if they exist
if landmarks:
for landmark in landmarks:
ax.scatter(landmark.x * image.shape[1], landmark.y * image.shape[0], color="red", s=10)
# Add height and area text
ax.text(10, 30, f"Predicted Gender: {gender}", color="yellow", fontsize=12, backgroundcolor='black')
ax.text(10, 60, f"Estimated Height: {real_height:.2f} cm", color="yellow", fontsize=12, backgroundcolor='black')
ax.text(10, 90, f"Estimated Torso Area: {real_torso_area:.2f} cm^2", color="yellow", fontsize=12, backgroundcolor='black')
# Add grid scale to background
ax.set_xticks(np.arange(0, image.shape[1], 50))
ax.set_yticks(np.arange(0, image.shape[0], 50))
ax.grid(color='white', linestyle='--', linewidth=0.5)
# Save annotated plot for output
plt.savefig("output_plot.png")
plt.close(fig)
return gender, real_height, real_torso_area, "output_plot.png"
except Exception as e:
print(f"Error encountered: {e}")
return f"An error occurred: {e}", None, None, None
# Define Gradio interface
image_input = gr.Image(label="Input Image")
text_output = gr.Textbox(label="Predicted Gender")
height_output = gr.Textbox(label="Estimated Height (cm)")
area_output = gr.Textbox(label="Estimated Torso Area (cm^2)")
image_output = gr.Image(label="Annotated Image with Measurements")
# Set up Gradio interface
gr_interface = gr.Interface(
fn=predict,
inputs=image_input,
outputs=[text_output, height_output, area_output, image_output],
title="Gender, Height, and Body Measurement Estimation",
description="Upload an image to predict gender, estimate height, and calculate 2D torso area with a scale overlay."
)
# Launch the Gradio app
gr_interface.launch()