import os
import cv2
import torch
import torchvision
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import pandas as pd
import psutil
import time

from smplx.lbs import batch_rodrigues

# Импортируйте ваши модули здесь
from configs import paths
from configs.pose2D_hrnet_config import get_pose2D_hrnet_cfg_defaults
from configs.poseMF_shapeGaussian_net_config import get_poseMF_shapeGaussian_cfg_defaults
from measure import MeasureBody
from models.canny_edge_detector import CannyEdgeDetector
from models.pose2D_hrnet import PoseHighResolutionNet
from models.poseMF_shapeGaussian_net import PoseMFShapeGaussianNet
from models.smpl_official import SMPL
from predict.predict_hrnet import predict_hrnet
from renderers.pytorch3d_textured_renderer import TexturedIUVRenderer
from utils.image_utils import batch_add_rgb_background, batch_crop_pytorch_affine, batch_crop_opencv_affine
from utils.label_conversions import convert_2Djoints_to_gaussian_heatmaps_torch
from utils.rigid_transform_utils import rot6d_to_rotmat, aa_rotate_translate_points_pytorch3d
from utils.sampling_utils import compute_vertex_uncertainties_by_poseMF_shapeGaussian_sampling, \
    joints2D_error_sorted_verts_sampling



class Inference:
    def __init__(self,
                 device: torch.device,
                 pose_shape_weights_path: str,
                 pose2D_hrnet_weights_path: str,
                 pose_shape_cfg_path: str = None,
                 gender: str = 'NEUTRAL',
                 height: int = 170,
                 mapping_config: dict = None,
                 scoring_config: dict = None,
                 cpu_threshold: int = 80,
                 memory_threshold: int = 80,
                 visualise_wh=512):
        """
        Инициализирует класс Inference с необходимыми моделями и конфигурациями.

        :param device: Устройство для вычислений (CPU или CUDA).
        :param pose_shape_weights_path: Путь к весам модели предсказания позы и формы.
        :param pose2D_hrnet_weights_path: Путь к весам модели HRNet для 2D-поз.
        :param pose_shape_cfg_path: Путь к конфигурационному файлу для модели предсказания позы и формы.
        :param gender: Пол модели SMPL ('NEUTRAL', 'MALE', 'FEMALE').
        :param height: Высота для нормализации замеров.
        :param mapping_config: Конфигурация для SizeTableMapper.
        :param scoring_config: Конфигурация для ScoreCalculator.
        :param cpu_threshold: Порог загрузки CPU для ResourceManager.
        :param memory_threshold: Порог загрузки памяти для ResourceManager.
        """
        self.device = device
        self.pose_shape_weights_path = pose_shape_weights_path
        self.pose2D_hrnet_weights_path = pose2D_hrnet_weights_path
        self.pose_shape_cfg_path = pose_shape_cfg_path
        self.gender = gender.upper()
        self.height = height
        self.mapping_config = mapping_config if mapping_config is not None else {}
        self.scoring_config = scoring_config if scoring_config is not None else {}
        self.cpu_threshold = cpu_threshold
        self.memory_threshold = memory_threshold
        self.visualise_wh=visualise_wh
        # Инициализация моделей и компонентов
        self.init_models()


        # Инициализация дополнительных компонентов
        self.size_table_mapper = self.SizeTableMapper(self.mapping_config)
        self.score_calculator = self.ScoreCalculator(self.scoring_config)
        # self.final_processor = self.FinalProcessor(
        #     measurement_processor=self.measurement_processor,
        #     size_table_mapper=self.size_table_mapper,
        #     score_calculator=self.score_calculator
        # )
        self.resource_manager = self.ResourceManager(cpu_threshold=self.cpu_threshold,
                                                     memory_threshold=self.memory_threshold)

    def init_models(self):
        """
        Инициализирует и загружает все необходимые модели, разделив процесс на мелкие методы.
        """
        self.load_configurations()
        self.load_object_detection_model()
        self.load_hrnet_model()
        self.load_edge_detector()
        self.load_smpl_model()
        self.load_pose_shape_model()
        self.init_visualizer()
        self.init_measurer()

    def load_configurations(self):
        """
        Загружает конфигурации моделей.
        """
        # Загрузка конфигураций
        self.pose2D_hrnet_cfg = get_pose2D_hrnet_cfg_defaults()
        self.pose_shape_cfg = get_poseMF_shapeGaussian_cfg_defaults()
        if self.pose_shape_cfg_path:
            self.pose_shape_cfg.merge_from_file(self.pose_shape_cfg_path)
            print('\nLoaded Distribution Predictor config from', self.pose_shape_cfg_path)
        else:
            print('\nUsing default Distribution Predictor config.')

    def load_object_detection_model(self):
        """
        Загружает модель обнаружения объектов, если изображения не обрезаны.
        """
        if not self.mapping_config.get('already_cropped_images', False):
            self.object_detect_model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True).to(self.device)
            self.object_detect_model.eval()
            print('\nLoaded Mask R-CNN model for object detection.')
        else:
            self.object_detect_model = None
            print('\nSkipping object detection model loading as images are already cropped.')

    def load_hrnet_model(self):
        """
        Загружает модель HRNet для обнаружения 2D суставов.
        """
        self.hrnet_model = PoseHighResolutionNet(self.pose2D_hrnet_cfg).to(self.device)
        hrnet_checkpoint = torch.load(self.pose2D_hrnet_weights_path, map_location=self.device)
        self.hrnet_model.load_state_dict(hrnet_checkpoint, strict=False)
        self.hrnet_model.eval()
        print('\nLoaded HRNet weights from', self.pose2D_hrnet_weights_path)

    def load_edge_detector(self):
        """
        Загружает детектор краев Canny.
        """
        self.edge_detect_model = CannyEdgeDetector(
            non_max_suppression=self.pose_shape_cfg.DATA.EDGE_NMS,
            gaussian_filter_std=self.pose_shape_cfg.DATA.EDGE_GAUSSIAN_STD,
            gaussian_filter_size=self.pose_shape_cfg.DATA.EDGE_GAUSSIAN_SIZE,
            threshold=self.pose_shape_cfg.DATA.EDGE_THRESHOLD
        ).to(self.device)
        self.edge_detect_model.eval()
        print('\nInitialized Canny Edge Detector.')

    def load_smpl_model(self):
        """
        Загружает модель SMPL.
        """
        print('\nUsing {} SMPL model with {} shape parameters.'.format(
            self.gender, str(self.pose_shape_cfg.MODEL.NUM_SMPL_BETAS)))
        self.smpl_model = SMPL(
            paths.SMPL,
            batch_size=1,
            gender=self.gender,
            num_betas=self.pose_shape_cfg.MODEL.NUM_SMPL_BETAS
        ).to(self.device)
        self.smpl_immediate_parents = self.smpl_model.parents.tolist()
        self.smpl_model.eval()
        print('\nLoaded SMPL model.')

    def load_pose_shape_model(self):
        """
        Загружает модель предсказания распределений позы и формы.
        """
        self.pose_shape_dist_model = PoseMFShapeGaussianNet(
            smpl_parents=self.smpl_immediate_parents,
            config=self.pose_shape_cfg
        ).to(self.device)
        pose_shape_checkpoint = torch.load(self.pose_shape_weights_path, map_location=self.device)
        self.pose_shape_dist_model.load_state_dict(pose_shape_checkpoint['best_model_state_dict'])
        self.pose_shape_dist_model.eval()
        print('\nLoaded Distribution Predictor weights from', self.pose_shape_weights_path)

    def init_visualizer(self):
        """
        Инициализирует визуализатор для отображения результатов.
        """
        self.body_vis_renderer = TexturedIUVRenderer(
            device=self.device,
            batch_size=1,
            img_wh=self.visualise_wh,  # Можно настроить
            projection_type='orthographic',
            render_rgb=True,
            bin_size=32
        )
        self.plain_texture = torch.ones(1, 1200, 800, 3, device=self.device).float() * 0.7
        self.lights_rgb_settings = {
            'location': torch.tensor([[0., -0.8, -2.0]], device=self.device, dtype=torch.float32),
            'ambient_color': 0.5 * torch.ones(1, 3, device=self.device, dtype=torch.float32),
            'diffuse_color': 0.3 * torch.ones(1, 3, device=self.device, dtype=torch.float32),
            'specular_color': torch.zeros(1, 3, device=self.device, dtype=torch.float32)
        }
        self.fixed_cam_t = torch.tensor([[0., -0.2, 2.5]], device=self.device)
        self.fixed_orthographic_scale = torch.tensor([[0.95, 0.95]], device=self.device)
        print('\nInitialized Visualizer.')

    def init_measurer(self):
        """
        Инициализирует измеритель, если требуется.
        """
        if self.pose_shape_cfg.get('get_measures', True):
            self.measurer = MeasureBody('smpl')
            print('\nInitialized MeasureBody.')
        else:
            self.measurer = None
            print('\nMeasureBody is disabled based on configuration.')

    def run_predict(self,
                   image_dir: str = None,
                   save_dir: str = None,
                   image: np.ndarray = None,
                   save_path: str = None,
                   object_detect_threshold: float = 0.75,
                   bbox_scale_factor: float = 1.2,
                   joints2Dvisib_threshold: float = 0.75,
                   visualise_uncropped: bool = True,
                   visualise_samples: bool = False,
                   get_measures: bool = True):
        """
        Выполняет предсказание на заданном наборе изображений или одном изображении.

        :param image_dir: Директория с входными изображениями.
        :param save_dir: Директория для сохранения результатов.
        :param image: Загруженное изображение в формате NumPy массива.
        :param save_path: Путь для сохранения результатов одного изображения.
        :param object_detect_threshold: Порог для обнаружения объектов.
        :param bbox_scale_factor: Масштабирование ограничивающей рамки.
        :param joints2Dvisib_threshold: Порог видимости 2D-соединений.
        :param visualise_wh: Ширина и высота визуализации.
        :param visualise_uncropped: Флаг для визуализации необрезанных изображений.
        :param visualise_samples: Флаг для визуализации выборок.
        :param get_measures: Флаг для получения замеров.
        """
        # Проверка ресурсов перед запуском
        self.resource_manager.wait_for_resources()
        if image_dir and save_dir:
            os.makedirs(save_dir, exist_ok=True)

            # Обработка каждого изображения в директории
            for image_fname in tqdm(sorted([f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))])):
                with torch.no_grad():
                    image_path = os.path.join(image_dir, image_fname)
                    self.process_single_image(
                        image_path=image_path,
                        save_dir=save_dir,
                        save_fname=image_fname,
                        object_detect_threshold=object_detect_threshold,
                        bbox_scale_factor=bbox_scale_factor,
                        joints2Dvisib_threshold=joints2Dvisib_threshold,
                        visualise_wh=self.visualise_wh,
                        visualise_uncropped=visualise_uncropped,
                        visualise_samples=visualise_samples,
                        get_measures=get_measures
                    )
        elif image is not None and save_path is not None:
            # Обработка одного изображения
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with torch.no_grad():
                self.process_single_image(
                    image=image,
                    save_dir=os.path.dirname(save_path),
                    save_fname=os.path.basename(save_path),
                    object_detect_threshold=object_detect_threshold,
                    bbox_scale_factor=bbox_scale_factor,
                    joints2Dvisib_threshold=joints2Dvisib_threshold,
                    visualise_wh=self.visualise_wh,
                    visualise_uncropped=visualise_uncropped,
                    visualise_samples=visualise_samples,
                    get_measures=get_measures
                )
        else:
            raise ValueError("Необходимо указать либо `image_dir` и `save_dir`, либо `image` и `save_path`.")

    def process_single_image(self,
                             image_path: str = None,
                             image: np.ndarray = None,
                             save_dir: str = None,
                             save_fname: str = None,
                             object_detect_threshold: float = 0.75,
                             bbox_scale_factor: float = 1.2,
                             visualise_wh:int =512,
                             joints2Dvisib_threshold: float = 0.75,
                             visualise_uncropped: bool = True,
                             visualise_samples: bool = False,
                             get_measures: bool = True):
        """
        Обрабатывает одно изображение и сохраняет результаты.

        :param image_path: Путь к входному изображению.
        :param image: Загруженное изображение в формате NumPy массива.
        :param save_dir: Директория для сохранения результатов.
        :param save_fname: Имя файла для сохранения результатов.
        :param object_detect_threshold: Порог для обнаружения объектов.
        :param bbox_scale_factor: Масштабирование ограничивающей рамки.
        :param joints2Dvisib_threshold: Порог видимости 2D-соединений.
        :param visualise_wh: Ширина и высота визуализации.
        :param visualise_uncropped: Флаг для визуализации необрезанных изображений.
        :param visualise_samples: Флаг для визуализации выборок.
        :param get_measures: Флаг для получения замеров.
        """
        # Загрузка изображения
        if image_path:
            image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
        elif image is not None:
            pass  # Изображение уже загружено
        else:
            raise ValueError("Необходимо предоставить либо `image_path`, либо `image`.")

        orig_image = image.copy()
        image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float().to(self.device) / 255.0

        # Предсказание ограничивающей рамки и 2D-соединений
        hrnet_output = predict_hrnet(
            hrnet_model=self.hrnet_model,
            hrnet_config=self.pose2D_hrnet_cfg,
            object_detect_model=self.object_detect_model,
            image=image_tensor,
            object_detect_threshold=object_detect_threshold,
            bbox_scale_factor=bbox_scale_factor
        )
        print(hrnet_output)

        # Преобразование предсказанных 2D суставов и изображения
        hrnet_input_centre = torch.tensor([[hrnet_output['cropped_image'].shape[1],
                                           hrnet_output['cropped_image'].shape[2]]],
                                         dtype=torch.float32,
                                         device=self.device) * 0.5
        hrnet_input_height = torch.tensor([hrnet_output['cropped_image'].shape[1]],
                                         dtype=torch.float32,
                                         device=self.device)
        cropped_for_proxy = batch_crop_pytorch_affine(
            input_wh=(self.pose2D_hrnet_cfg.MODEL.IMAGE_SIZE[0], self.pose2D_hrnet_cfg.MODEL.IMAGE_SIZE[1]),
            output_wh=(self.pose_shape_cfg.DATA.PROXY_REP_SIZE, self.pose_shape_cfg.DATA.PROXY_REP_SIZE),
            num_to_crop=1,
            device=self.device,
            joints2D=hrnet_output['joints2D'][None, :, :],
            rgb=hrnet_output['cropped_image'][None, :, :, :],
            bbox_centres=hrnet_input_centre,
            bbox_heights=hrnet_input_height,
            bbox_widths=hrnet_input_height,
            orig_scale_factor=1.0
        )

        # Создание прокси-представления с использованием обнаружения краев и горячих карт 2D-соединений
        edge_detector_output = self.edge_detect_model(cropped_for_proxy['rgb'])
        proxy_rep_img = edge_detector_output['thresholded_thin_edges'] if self.pose_shape_cfg.DATA.EDGE_NMS else \
            edge_detector_output['thresholded_grad_magnitude']
        proxy_rep_heatmaps = convert_2Djoints_to_gaussian_heatmaps_torch(
            joints2D=cropped_for_proxy['joints2D'],
            img_wh=self.pose_shape_cfg.DATA.PROXY_REP_SIZE,
            std=self.pose_shape_cfg.DATA.HEATMAP_GAUSSIAN_STD
        )
        hrnet_joints2Dvisib = hrnet_output['joints2Dconfs'] > joints2Dvisib_threshold
        hrnet_joints2Dvisib[[0, 1, 2, 3, 4, 5, 6, 11, 12]] = True  # Фиксируем видимость некоторых суставов
        proxy_rep_heatmaps = proxy_rep_heatmaps * hrnet_joints2Dvisib[None, :, None, None]
        proxy_rep_input = torch.cat([proxy_rep_img, proxy_rep_heatmaps], dim=1).float()  # (1, 18, img_wh, img_wh)

        # Предсказание распределений позы и формы
        pred_pose_F, pred_pose_U, pred_pose_S, pred_pose_V, pred_pose_rotmats_mode, \
            pred_shape_dist, pred_glob, pred_cam_wp = self.pose_shape_dist_model(proxy_rep_input)

        # Обработка глобальной ориентации
        if pred_glob.shape[-1] == 3:
            pred_glob_rotmats = batch_rodrigues(pred_glob)  # (1, 3, 3)
        elif pred_glob.shape[-1] == 6:
            pred_glob_rotmats = rot6d_to_rotmat(pred_glob)  # (1, 3, 3)
        else:
            raise ValueError("Unexpected shape for pred_glob: {}".format(pred_glob.shape))

        # Получение SMPL параметров
        pred_smpl_output_mode = self.smpl_model(
            body_pose=pred_pose_rotmats_mode,
            global_orient=pred_glob_rotmats.unsqueeze(1),
            betas=pred_shape_dist.loc,
            pose2rot=False
        )

        # Получение параметров формы (betas)
        measurments = None
        if get_measures and self.measurer:
            betas = pred_shape_dist.loc.cpu().numpy()
            betas = torch.tensor(betas)

            self.measurer.from_body_model(shape=betas, gender=self.gender)
            all_measurement_names = self.measurer.all_possible_measurements
            self.measurer.measure(all_measurement_names)
            new_height = self.height
            self.measurer.height_normalize_measurements(new_height)
            measurments = self.measurer.height_normalized_measurements

        pred_vertices_mode = pred_smpl_output_mode.vertices  # (1, 6890, 3)
        # Поворот вершин для правильной проекции
        pred_vertices_mode = aa_rotate_translate_points_pytorch3d(
            points=pred_vertices_mode,
            axes=torch.tensor([1., 0., 0.], device=self.device),
            angles=np.pi,
            translations=torch.zeros(3, device=self.device)
        )
        pred_vertices_rot90_mode = aa_rotate_translate_points_pytorch3d(
            points=pred_vertices_mode,
            axes=torch.tensor([0., 1., 0.], device=self.device),
            angles=-np.pi / 2.,
            translations=torch.zeros(3, device=self.device)
        )
        pred_vertices_rot180_mode = aa_rotate_translate_points_pytorch3d(
            points=pred_vertices_rot90_mode,
            axes=torch.tensor([0., 1., 0.], device=self.device),
            angles=-np.pi / 2.,
            translations=torch.zeros(3, device=self.device)
        )
        pred_vertices_rot270_mode = aa_rotate_translate_points_pytorch3d(
            points=pred_vertices_rot180_mode,
            axes=torch.tensor([0., 1., 0.], device=self.device),
            angles=-np.pi / 2.,
            translations=torch.zeros(3, device=self.device)
        )

        pred_reposed_smpl_output_mean = self.smpl_model(
            betas=pred_shape_dist.loc
        )
        pred_reposed_vertices_mean = pred_reposed_smpl_output_mean.vertices  # (1, 6890, 3)
        pred_reposed_vertices_flipped_mean = aa_rotate_translate_points_pytorch3d(
            points=pred_reposed_vertices_mean,
            axes=torch.tensor([1., 0., 0.], device=self.device),
            angles=np.pi,
            translations=torch.zeros(3, device=self.device)
        )
        pred_reposed_vertices_rot90_mean = aa_rotate_translate_points_pytorch3d(
            points=pred_reposed_vertices_flipped_mean,
            axes=torch.tensor([0., 1., 0.], device=self.device),
            angles=-np.pi / 2.,
            translations=torch.zeros(3, device=self.device)
        )

        # Визуализация
        orthographic_scale = pred_cam_wp[:, [0, 0]]
        cam_t = torch.cat([
            pred_cam_wp[:, 1:],
            torch.ones(pred_cam_wp.shape[0], 1, device=self.device).float() * 2.5
        ], dim=-1)

        # Оценка неопределенности вершин путем семплинга
        per_vertex_3Dvar, pred_vertices_samples, pred_joints_samples = \
            compute_vertex_uncertainties_by_poseMF_shapeGaussian_sampling(
                pose_U=pred_pose_U,
                pose_S=pred_pose_S,
                pose_V=pred_pose_V,
                shape_distribution=pred_shape_dist,
                glob_rotmats=pred_glob_rotmats,
                num_samples=50,
                smpl_model=self.smpl_model,
                use_mean_shape=True
            )

        if visualise_samples:
            num_samples = 8
            pred_vertices_samples = joints2D_error_sorted_verts_sampling(
                pred_vertices_samples=pred_vertices_samples,
                pred_joints_samples=pred_joints_samples,
                input_joints2D_heatmaps=proxy_rep_input[:, 1:, :, :],
                pred_cam_wp=pred_cam_wp
            )[:num_samples, :, :]  # (8, 6890, 3)
            pred_vertices_samples = aa_rotate_translate_points_pytorch3d(
                points=pred_vertices_samples,
                axes=torch.tensor([1., 0., 0.], device=self.device),
                angles=np.pi,
                translations=torch.zeros(3, device=self.device)
            )
            pred_vertices_rot90_samples = aa_rotate_translate_points_pytorch3d(
                points=pred_vertices_samples,
                axes=torch.tensor([0., 1., 0.], device=self.device),
                angles=-np.pi / 2.,
                translations=torch.zeros(3, device=self.device)
            )
            pred_vertices_samples = torch.cat([pred_vertices_mode, pred_vertices_samples], dim=0)  # (9, 6890, 3)
            pred_vertices_rot90_samples = torch.cat([pred_vertices_rot90_mode, pred_vertices_rot90_samples], dim=0)  # (9, 6890, 3)

        # Генерация цветовой карты неопределенности вершин
        vertex_var_norm = plt.Normalize(vmin=0.0, vmax=0.2, clip=True)
        vertex_var_colours = plt.cm.turbo(vertex_var_norm(per_vertex_3Dvar.cpu().detach().numpy()))[:, :3]
        vertex_var_colours = torch.from_numpy(vertex_var_colours[None, :, :]).to(self.device).float()


        # Рендеринг визуализаций
        body_vis_output = self.body_vis_renderer(
            vertices=pred_vertices_mode,
            cam_t=cam_t,
            orthographic_scale=orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings,
            verts_features=vertex_var_colours
        )
        cropped_for_proxy_rgb = torch.nn.functional.interpolate(
            cropped_for_proxy['rgb'],
            size=(visualise_wh, visualise_wh),
            mode='bilinear',
            align_corners=False
        )
        body_vis_rgb = batch_add_rgb_background(
            backgrounds=cropped_for_proxy_rgb,
            rgb=body_vis_output['rgb_images'].permute(0, 3, 1, 2).contiguous(),
            seg=body_vis_output['iuv_images'][:, :, :, 0].round()
        )
        body_vis_rgb = body_vis_rgb.cpu().detach().numpy()[0].transpose(1, 2, 0)

        # Вращенные визуализации
        body_vis_rgb_rot90 = self.body_vis_renderer(
            vertices=pred_vertices_rot90_mode,
            cam_t=self.fixed_cam_t,
            orthographic_scale=self.fixed_orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings,
            verts_features=vertex_var_colours
        )['rgb_images'].cpu().detach().numpy()[0]

        body_vis_rgb_rot180 = self.body_vis_renderer(
            vertices=pred_vertices_rot180_mode,
            cam_t=self.fixed_cam_t,
            orthographic_scale=self.fixed_orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings,
            verts_features=vertex_var_colours
        )['rgb_images'].cpu().detach().numpy()[0]

        body_vis_rgb_rot270 = self.body_vis_renderer(
            vertices=pred_vertices_rot270_mode,
            textures=self.plain_texture,
            cam_t=self.fixed_cam_t,
            orthographic_scale=self.fixed_orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings,
            verts_features=vertex_var_colours
        )['rgb_images'].cpu().detach().numpy()[0]

        # Визуализация ре-позированных вершин
        reposed_body_vis_rgb = self.body_vis_renderer(
            vertices=pred_reposed_vertices_flipped_mean,
            textures=self.plain_texture,
            cam_t=self.fixed_cam_t,
            orthographic_scale=self.fixed_orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings
        )['rgb_images'].cpu().detach().numpy()[0]

        reposed_body_vis_rgb_rot90 = self.body_vis_renderer(
            vertices=pred_reposed_vertices_rot90_mean,
            textures=self.plain_texture,
            cam_t=self.fixed_cam_t,
            orthographic_scale=self.fixed_orthographic_scale,
            lights_rgb_settings=self.lights_rgb_settings
        )['rgb_images'].cpu().detach().numpy()[0]

        # Объединение всех визуализаций в сетку
        combined_vis_rows = 2
        combined_vis_cols = 4  # Увеличиваем количество колонок на 1
        combined_vis_fig = np.zeros((combined_vis_rows * visualise_wh,
                                     combined_vis_cols * visualise_wh, 3),
                                    dtype=body_vis_rgb.dtype)
        combined_vis_fig.fill(254)

        # Вставка обрезанного исходного изображения
        combined_vis_fig[:visualise_wh, :visualise_wh] = cropped_for_proxy_rgb.cpu().detach().numpy()[0].transpose(1, 2, 0)

        # Вставка текстового наложения с замерами
        text_overlay = np.zeros((visualise_wh, visualise_wh, 3), dtype=np.uint8)  # Пустое изображение для текста
        y_offset = 20
        if measurments:
            for measurement_name, value in measurments.items():
                cv2.putText(text_overlay, f"{measurement_name}: {value:.2f}", (10, y_offset),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), lineType=cv2.LINE_AA)
                y_offset += 30

        combined_vis_fig[visualise_wh:2 * visualise_wh, :visualise_wh] = text_overlay

        # Вставка рендеринговых изображений
        combined_vis_fig[:visualise_wh, visualise_wh:2 * visualise_wh] = body_vis_rgb
        combined_vis_fig[visualise_wh:2 * visualise_wh, visualise_wh:2 * visualise_wh] = body_vis_rgb_rot90
        combined_vis_fig[:visualise_wh, 2 * visualise_wh:3 * visualise_wh] = body_vis_rgb_rot180
        combined_vis_fig[visualise_wh:2 * visualise_wh, 2 * visualise_wh:3 * visualise_wh] = body_vis_rgb_rot270

        combined_vis_fig[:visualise_wh, 3 * visualise_wh:4 * visualise_wh] = reposed_body_vis_rgb
        combined_vis_fig[visualise_wh:2 * visualise_wh, 3 * visualise_wh:4 * visualise_wh] = reposed_body_vis_rgb_rot90

        # Сохранение визуализации
        vis_save_path = os.path.join(save_dir, save_fname)
        cv2.imwrite(vis_save_path, combined_vis_fig[:, :, ::-1] * 255)
        print(f"Saved visualization to {vis_save_path}")

        if visualise_uncropped:
            # Визуализация необрезанного изображения путем проекции 3D тела на оригинальное изображение
            rgb_to_uncrop = body_vis_output['rgb_images'].permute(0, 3, 1, 2).contiguous().cpu().detach().numpy()
            iuv_to_uncrop = body_vis_output['iuv_images'].permute(0, 3, 1, 2).contiguous().cpu().detach().numpy()
            bbox_centres = hrnet_output['bbox_centre'][None].cpu().detach().numpy()
            bbox_whs = torch.max(hrnet_output['bbox_height'], hrnet_output['bbox_width'])[None].cpu().detach().numpy()
            bbox_whs *= bbox_scale_factor
            uncropped_for_visualise = batch_crop_opencv_affine(
                output_wh=(visualise_wh, visualise_wh),
                num_to_crop=1,
                rgb=rgb_to_uncrop,
                iuv=iuv_to_uncrop,
                bbox_centres=bbox_centres,
                bbox_whs=bbox_whs,
                uncrop=True,
                uncrop_wh=(orig_image.shape[1], orig_image.shape[0])
            )
            uncropped_rgb = uncropped_for_visualise['rgb'][0].transpose(1, 2, 0) * 255
            uncropped_seg = uncropped_for_visualise['iuv'][0, 0, :, :]
            background_pixels = uncropped_seg[:, :, None] == 0  # Тело pixels > 0
            uncropped_rgb_with_background = uncropped_rgb * (~background_pixels) + \
                                            orig_image * background_pixels

            uncropped_vis_save_path = os.path.splitext(vis_save_path)[0] + '_uncrop.png'
            cv2.imwrite(uncropped_vis_save_path, uncropped_rgb_with_background[:, :, ::-1])
            print(f"Saved uncropped visualization to {uncropped_vis_save_path}")

        if visualise_samples:
            samples_rows = 3
            samples_cols = 6
            samples_fig = np.zeros((samples_rows * visualise_wh, samples_cols * visualise_wh, 3),
                                   dtype=body_vis_rgb.dtype)
            for i in range(num_samples + 1):
                body_vis_output_sample = self.body_vis_renderer(
                    vertices=pred_vertices_samples[[i]],
                    textures=self.plain_texture,
                    cam_t=cam_t,
                    orthographic_scale=orthographic_scale,
                    lights_rgb_settings=self.lights_rgb_settings
                )
                body_vis_rgb_sample = batch_add_rgb_background(
                    backgrounds=cropped_for_proxy_rgb,
                    rgb=body_vis_output_sample['rgb_images'].permute(0, 3, 1, 2).contiguous(),
                    seg=body_vis_output_sample['iuv_images'][:, :, :, 0].round()
                )
                body_vis_rgb_sample = body_vis_rgb_sample.cpu().detach().numpy()[0].transpose(1, 2, 0)

                body_vis_rgb_rot90_sample = self.body_vis_renderer(
                    vertices=pred_vertices_rot90_samples[[i]],
                    textures=self.plain_texture,
                    cam_t=self.fixed_cam_t,
                    orthographic_scale=self.fixed_orthographic_scale,
                    lights_rgb_settings=self.lights_rgb_settings
                )['rgb_images'].cpu().detach().numpy()[0]

                row = (2 * i) // samples_cols
                col = (2 * i) % samples_cols
                samples_fig[row * visualise_wh:(row + 1) * visualise_wh,
                            col * visualise_wh:(col + 1) * visualise_wh] = body_vis_rgb_sample

                row = (2 * i + 1) // samples_cols
                col = (2 * i + 1) % samples_cols
                samples_fig[row * visualise_wh:(row + 1) * visualise_wh,
                            col * visualise_wh:(col + 1) * visualise_wh] = body_vis_rgb_rot90_sample

            samples_fig_save_path = os.path.splitext(vis_save_path)[0] + '_samples.png'
            cv2.imwrite(samples_fig_save_path, samples_fig[:, :, ::-1] * 255)
            print(f"Saved samples visualization to {samples_fig_save_path}")

    class SizeTableMapper:
        def __init__(self, mapping_config: dict):
            """
            Инициализирует маппер размерной таблицы.

            :param mapping_config: Конфигурация для маппинга.
            """
            self.mapping_config = mapping_config
            self.mapping_table = None

        def generate_mapping(self, measurements: dict):
            """
            Генерирует маппинг размерной таблицы на основе замеров.

            :param measurements: Замеры тела.
            :return: DataFrame с маппингом размерной таблицы.
            """
            # Реализуйте логику формирования маппинга
            # Здесь представлен простой пример
            self.mapping_table = pd.DataFrame([measurements])
            # Примените необходимые преобразования согласно конфигурации
            # Например, нормализация, преобразование единиц и т.д.
            # Пример:
            # for key, transform in self.mapping_config.items():
            #     self.mapping_table[key] = self.mapping_table[key].apply(transform)
            return self.mapping_table

        def save_mapping(self, save_path: str):
            """
            Сохраняет маппинг размерной таблицы в файл.

            :param save_path: Путь для сохранения CSV-файла.
            """
            if self.mapping_table is not None:
                self.mapping_table.to_csv(save_path, index=False)
                print(f"Mapping table saved to {save_path}")

    class ScoreCalculator:
        def __init__(self, scoring_config: dict):
            """
            Инициализирует калькулятор скоров.

            :param scoring_config: Конфигурация для расчета скоров.
            """
            self.scoring_config = scoring_config

        def calculate_score(self, measurements: dict):
            """
            Рассчитывает скоры на основе замеров.

            :param measurements: Замеры тела.
            :return: Словарь со скоринговыми значениями.
            """
            scores = {}
            for key, value in measurements.items():
                if key in self.scoring_config:
                    min_val = self.scoring_config[key].get('min', 0)
                    max_val = self.scoring_config[key].get('max', 1)
                    norm_score = self.normalize(value, min_val, max_val)
                    scores[key] = norm_score
                else:
                    # Если конфигурация для ключа отсутствует, пропускаем
                    print(f"No scoring config for {key}, skipping.")
            return scores

        @staticmethod
        def normalize(value: float, min_val: float, max_val: float):
            """
            Нормализует значение в диапазоне от 0 до 1.

            :param value: Исходное значение.
            :param min_val: Минимальное значение.
            :param max_val: Максимальное значение.
            :return: Нормализованное значение.
            """
            if max_val - min_val == 0:
                return 0
            return (value - min_val) / (max_val - min_val)

    class FinalProcessor:
        def __init__(self, measurement_processor, size_table_mapper, score_calculator):
            """
            Инициализирует итоговый процессор.

            :param measurement_processor: Экземпляр MeasurementProcessor.
            :param size_table_mapper: Экземпляр SizeTableMapper.
            :param score_calculator: Экземпляр ScoreCalculator.
            """
            self.measurement_processor = measurement_processor
            self.size_table_mapper = size_table_mapper
            self.score_calculator = score_calculator

        def execute(self, image_path: str, save_path: str, mapping_save_path: str):
            """
            Выполняет полный процесс обработки изображения.

            :param image_path: Путь к входному изображению.
            :param save_path: Путь для сохранения промежуточных результатов.
            :param mapping_save_path: Путь для сохранения маппинга.
            :return: Финальный результат обработки.
            """
            # Шаг 1: Обработка замеров по фото
            self.measurement_processor.process_image(image_path, save_path)

            # Шаг 2: Формирование маппинга размерной таблицы
            measurements = self.get_measurements(save_path)
            mapping = self.size_table_mapper.generate_mapping(measurements)
            self.size_table_mapper.save_mapping(mapping_save_path)

            # Шаг 3: Расчет скоров
            scores = self.score_calculator.calculate_score(measurements)

            # Шаг 4: Итоговый результат
            final_result = {
                'measurements': measurements,
                'mapping': mapping,
                'scores': scores
            }
            return final_result

        def get_measurements(self, save_path: str):
            """
            Извлекает замеры из сохраненных данных.

            :param save_path: Путь к папке с сохраненными данными.
            :return: Словарь с замерами.
            """
            # Реализуйте логику извлечения замеров
            # Например, загрузка из файла или базы данных
            measurements = {}
            # Пример заполнения:
            # Предполагается, что save_path содержит JSON или другой формат с замерами
            # Здесь представлен заглушка
            measurements['height'] = self.measurement_processor.height
            # Добавьте остальные замеры по необходимости
            return measurements

    class ResourceManager:
        def __init__(self, cpu_threshold: int = 80, memory_threshold: int = 80):
            """
            Инициализирует менеджер ресурсов.

            :param cpu_threshold: Порог загрузки CPU (в процентах).
            :param memory_threshold: Порог загрузки памяти (в процентах).
            """
            self.cpu_threshold = cpu_threshold
            self.memory_threshold = memory_threshold

        def check_resources(self):
            """
            Проверяет текущую загрузку ресурсов.

            :return: True, если ресурсы в пределах допустимых значений, иначе False.
            """
            cpu_usage = psutil.cpu_percent(interval=1)
            memory_usage = psutil.virtual_memory().percent
            print(f"CPU Usage: {cpu_usage}%, Memory Usage: {memory_usage}%")
            if cpu_usage > self.cpu_threshold or memory_usage > self.memory_threshold:
                return False
            return True

        def wait_for_resources(self):
            """
            Ожидает, пока загрузка ресурсов не станет ниже пороговых значений.
            """
            while not self.check_resources():
                print("High resource usage detected. Waiting for resources to be available...")
                time.sleep(5)  # Ждем 5 секунд перед следующей проверкой
            print("Resources are available. Starting prediction.")

    def process_image(self, image_path: str = None, save_dir: str = None):
        """
        Обрабатывает одно изображение и сохраняет результаты.

        :param image_path: Путь к входному изображению.
        :param save_dir: Директория для сохранения результатов.
        """
        if not image_path or not save_dir:
            raise ValueError("Необходимо предоставить `image_path` и `save_dir`.")

        # Создаем директорию для сохранения, если не существует
        individual_save_dir = os.path.join(save_dir, os.path.splitext(os.path.basename(image_path))[0])
        os.makedirs(individual_save_dir, exist_ok=True)

        # Выполняем предсказание
        self.run_predict(
            image_dir=os.path.dirname(image_path),
            save_dir=individual_save_dir
        )


# Пример использования класса Inference
def main():
    import sys
    sys.path.append('/content/drive/Othercomputers/Mac/sizer3D')
    # from runner import Inference  # Предполагается, что класс Inference находится в runner.py

    import os
    import torch

    root = '/content/drive/Othercomputers/Mac/sizer3D'
    in_folders = ['demo_my']  # Можно добавить другие папки

    # Установите параметры
    for in_folder in in_folders:
        image_dir = f'{root}/{in_folder}/'  # Путь к директории с изображениями
        save_dir = f'{root}/{in_folder}_output/'  # Путь к директории для сохранения результатов
        pose_shape_weights_path = f'{root}/model_files/poseMF_shapeGaussian_net_weights.tar'  # Путь к весам модели
        pose2D_hrnet_weights_path = f'{root}/model_files/pose_hrnet_w48_384x288.pth'  # Путь к весам HRNet
        pose_shape_cfg_path = None  # Путь к конфигурации модели (если есть)
        already_cropped_images = False  # Укажите, если изображения уже обрезаны
        visualise_samples = True  # Визуализировать образцы
        visualise_uncropped = False  # Визуализировать не обрезанные изображения
        joints2Dvisib_threshold = 0.75  # Порог видимости 2D суставов
        gender = 'NEUTRAL'  # Пол для SMPL модели NEUTRAL, MALE, FEMALE


        # Установите устройство (GPU или CPU)
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print('\nDevice: {}'.format(device))

        # Создайте директорию для сохранения, если она не существует
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        # Инициализация класса Inference
        inference = Inference(
            device=device,
            pose_shape_weights_path=pose_shape_weights_path,
            pose2D_hrnet_weights_path=pose2D_hrnet_weights_path,
            pose_shape_cfg_path=pose_shape_cfg_path,
            gender=gender,
            height=170,  # Можно настроить
            mapping_config={
                # Ваша конфигурация для SizeTableMapper, если есть
            },
            scoring_config={
                'height': {'min': 140, 'max': 200},
                # Добавьте остальные параметры замеров и их диапазоны
            },
            cpu_threshold=80,
            memory_threshold=80,
            visualise_wh=512
        )

        # Запустите функцию предсказания
        inference.run_predict(
            image_dir=image_dir,
            save_dir=save_dir,
            object_detect_threshold=joints2Dvisib_threshold,
            bbox_scale_factor=1.2,
            joints2Dvisib_threshold=joints2Dvisib_threshold,
            visualise_uncropped=visualise_uncropped,
            visualise_samples=visualise_samples,
            get_measures=True
        )



if __name__ == "__main__":
    main()
отрефакторь