single_person_processor.py

# single_person_processor.py

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # Suppress TensorFlow warnings
import io
import logging
import numpy as np
import tensorflow as tf
from rembg import new_session, remove
from PIL import Image
from typing import Tuple, Dict
from tensorflow.keras.preprocessing.image import img_to_array

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class ImagePreprocessor:
    """Class to preprocess images by removing background."""

    def __init__(self, img_size: Tuple[int, int] = (128, 128)):
        """
        Initialize the image preprocessor.

        Parameters:
        - img_size: Target size for images (default is (128, 128)).
        """
        self.img_size = img_size
        self.session = new_session()  # Create a new session for rembg
        logging.info("Image processor initialized.")

    def process_single_person(self, front_path: str, side_path: str) -> Tuple[np.ndarray, np.ndarray]:
        """
        Process two images (front and side views).

        Parameters:
        - front_path: Path to the front image.
        - side_path: Path to the side image.

        Returns:
        - A tuple containing processed front and side images as NumPy arrays.
        """
        return (
            self._process_image(front_path),
            self._process_image(side_path)
        )

    def _process_image(self, image_path: str) -> np.ndarray:
        """
        Remove background and process a single image.

        Parameters:
        - image_path: Path to the image file.

        Returns:
        - Processed image as a NumPy array.
        """
        try:
            with open(image_path, "rb") as f:
                img_bytes = f.read()

            # Remove background using rembg
            bg_removed = Image.open(io.BytesIO(remove(img_bytes, session=self.session)))

            # Convert person to white and background to black
            white_img = Image.new("RGBA", bg_removed.size, (255, 255, 255, 255))
            white_img.putalpha(bg_removed.getchannel('A'))
            final_img = Image.alpha_composite(
                Image.new("RGBA", white_img.size, (0, 0, 0, 255)),
                white_img
            )

            # Prepare the image for the model
            return self._prepare_image(final_img)

        except Exception as e:
            logging.error(f"Failed to process {image_path}: {str(e)}")
            raise

    def _prepare_image(self, image: Image.Image) -> np.ndarray:
        """
        Convert the image to the model's input format.

        Parameters:
        - image: PIL Image object.

        Returns:
        - Image as a normalized NumPy array.
        """
        return img_to_array(
            image.convert('L').resize(self.img_size)
        ).astype(np.float32) / 255.0


class SinglePersonPredictor:
    """Class to predict body measurements and clothing sizes."""

    MEASUREMENT_INDICES = {
        'ankle': 0, 'arm-length': 1, 'bicep': 2, 'calf': 3,
        'chest': 4, 'forearm': 5, 'height': 6, 'hip': 7,
        'leg-length': 8, 'shoulder-breadth': 9,
        'shoulder-to-crotch': 10, 'thigh': 11, 'waist': 12, 'wrist': 13
    }

    SIZE_CHARTS = {
        'male': {
            'tshirt': [
                (97, 42, 'S'), (104, 45, 'M'), (112, 48, 'L'),
                (120, 51, 'XL'), (128, 54, 'XXL'), (136, 57, 'XXXL')
            ],
            'pants': [
                (76, 102, 30), (81, 107, 32), (86, 112, 34),
                (91, 117, 36), (97, 122, 38), (102, 127, 40), (107, 132, 42)
            ]
        },
        'female': {
            'tshirt': [
                (89, 38, 'S'), (96, 41, 'M'), (104, 44, 'L'),
                (112, 47, 'XL'), (120, 50, 'XXL'), (128, 53, 'XXXL')
            ],
            'pants': [
                (66, 92, 26), (71, 97, 28), (76, 102, 30),
                (81, 107, 32), (86, 112, 34), (91, 117, 36), (97, 122, 38)
            ]
        }
    }

    def __init__(self, model_path: str = 'best_model.keras'):
        """
        Initialize the predictor.

        Parameters:
        - model_path: Path to the trained model.
        """
        self.model = tf.keras.models.load_model(model_path)
        self.preprocessor = ImagePreprocessor()
        logging.info("Model loaded successfully.")

    def predict_measurements(self,
                             front_img_path: str,
                             side_img_path: str,
                             gender: int,
                             height_cm: float,
                             weight_kg: float,
                             apparel_type: str = "all") -> Dict:
        """
        Perform predictions for a single person and calculate clothing sizes.

        Parameters:
        - front_img_path: Path to the front image.
        - side_img_path: Path to the side image.
        - gender: 0 for male, 1 for female.
        - height_cm: Height in centimeters.
        - weight_kg: Weight in kilograms.
        - apparel_type: Specify "tshirt", "pants", or "all".

        Returns:
        - Dictionary containing predicted measurements and clothing sizes.
        """
        try:
            # Validate inputs
            if gender not in (0, 1):
                raise ValueError("Gender must be 0 (male) or 1 (female).")
            if not 100 <= height_cm <= 250:
                raise ValueError("Height must be between 100-250 cm.")
            if not 30 <= weight_kg <= 300:
                raise ValueError("Weight must be between 30-300 kg.")
            if apparel_type not in ["tshirt", "pants", "all"]:
                raise ValueError("Apparel type must be 'tshirt', 'pants', or 'all'.")

            # Process images
            front_arr, side_arr = self.preprocessor.process_single_person(front_img_path, side_img_path)
            meta_arr = np.array([[gender, height_cm, weight_kg]], dtype=np.float32)

            # Make predictions
            prediction = self.model.predict([
                np.expand_dims(front_arr, axis=0),
                np.expand_dims(side_arr, axis=0),
                meta_arr
            ])

            # Convert predictions to dictionary
            measurements = {
                name: round(float(prediction[0][idx]), 2)
                for name, idx in self.MEASUREMENT_INDICES.items()
            }

            result = {}

            if apparel_type == "tshirt":
                result["tshirt_size"] = self.calculate_tshirt_size(gender, measurements, weight_kg)

            elif apparel_type == "pants":
                result["pants_size"] = self.calculate_pants_size(gender, measurements, weight_kg)

            elif apparel_type == "all":
                result["body_measurements"] = measurements
                tshirt_size, pants_size = self.calculate_apparel_size(gender, measurements, weight_kg)
                result["tshirt_size"] = tshirt_size
                result["pants_size"] = pants_size

            return result

        except Exception as e:
            logging.error(f"Prediction failed: {str(e)}")
            raise

    def calculate_tshirt_size(self, gender: int, measurements: Dict, weight: float) -> str:
        """
        Calculate t-shirt size based on chest and shoulder-breadth measurements.

        Parameters:
        - gender: 0 for male, 1 for female.
        - measurements: Dictionary of body measurements.
        - weight: Weight in kilograms.

        Returns:
        - T-Shirt size as a string.
        """
        gender_str = 'male' if gender == 0 else 'female'
        chart = self.SIZE_CHARTS[gender_str]['tshirt']

        chest = measurements['chest']
        shoulder_breadth = measurements['shoulder-breadth']

        base_size = next(
            (size for max_chest, max_shoulder, size in chart
             if chest <= max_chest and shoulder_breadth <= max_shoulder),
            'XXXL'
        )

        if weight > 95 and gender == 0 or weight > 80 and gender == 1:
            base_size = 'XXL'

        if measurements['height'] > 180:
            base_size = f"Tall {base_size}"

        return base_size

    def calculate_pants_size(self, gender: int, measurements: Dict, weight: float) -> int:
        """
        Calculate pants size based on waist and hip measurements.

        Parameters:
        - gender: 0 for male, 1 for female.
        - measurements: Dictionary of body measurements.
        - weight: Weight in kilograms.

        Returns:
        - Pants size as an integer.
        """
        gender_str = 'male' if gender == 0 else 'female'
        chart = self.SIZE_CHARTS[gender_str]['pants']

        waist = measurements['waist']
        hip = measurements['hip']

        base_size = next(
            (size for max_waist, max_hip, size in chart
             if waist <= max_waist and hip <= max_hip),
            chart[-1][2]
        )

        if measurements['height'] > 180 and gender == 0:
            base_size += 2 if base_size < 40 else 0

        return base_size

    def calculate_apparel_size(self, gender: int, measurements: Dict, weight: float) -> Tuple[str, int]:
        """
        Calculate both t-shirt and pants sizes.

        Parameters:
        - gender: 0 for male, 1 for female.
        - measurements: Dictionary of body measurements.
        - weight: Weight in kilograms.

        Returns:
        - A tuple containing t-shirt size and pants size.
        """
        tshirt_size = self.calculate_tshirt_size(gender, measurements, weight)
        pants_size = self.calculate_pants_size(gender, measurements, weight)
        return tshirt_size, pants_size