Реализация AI-виртуальной примерки одежды

Разработка AI-системы виртуальной примерки одежды

Virtual try-on позволяет пользователям увидеть, как будет выглядеть одежда на них лично: загружаешь фото, выбираешь товар — система генерирует реалистичный результат. Снижает returns на 20–40% в fashion e-commerce, увеличивает конверсию на 10–25%.

Технологический стек

IDM-VTON (2024) — текущий SOTA для виртуальной примерки одежды:

import torch
from diffusers import AutoPipelineForInpainting
from transformers import AutoProcessor, CLIPVisionModelWithProjection
import numpy as np
from PIL import Image
import io

class VirtualTryOnService:
    def __init__(self):
        # IDM-VTON основан на SDXL inpainting + специализированный encoder
        self.pipeline = self._load_idm_vton()
        self.parsing_model = self._load_human_parsing()  # SCHP / CIHP
        self.pose_estimator = self._load_pose_estimator()  # OpenPose / DWPose

    def _load_idm_vton(self):
        from idm_vton import TryOnPipeline
        return TryOnPipeline.from_pretrained(
            "yisol/IDM-VTON",
            torch_dtype=torch.float16
        ).to("cuda")

    def try_on(
        self,
        person_image: bytes,
        garment_image: bytes,
        garment_description: str = "",
        seed: int = 42,
        num_steps: int = 30
    ) -> bytes:
        person_pil = Image.open(io.BytesIO(person_image)).convert("RGB")
        garment_pil = Image.open(io.BytesIO(garment_image)).convert("RGB")

        # Парсинг тела: определяем зону для примерки
        person_parse = self.parsing_model(person_pil)
        pose_map = self.pose_estimator(person_pil)

        result = self.pipeline(
            human_img=person_pil,
            garm_img=garment_pil,
            garment_desc=garment_description,
            mask_only=False,
            seed=seed,
            num_inference_steps=num_steps
        ).images[0]

        buf = io.BytesIO()
        result.save(buf, format="PNG")
        return buf.getvalue()

Human Parsing (сегментация тела)

from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
import torch

class HumanBodyParser:
    LABELS = {
        0: "background", 1: "hat", 2: "hair", 4: "upper-clothes",
        5: "skirt", 6: "pants", 7: "dress", 9: "belt",
        10: "left-shoe", 11: "right-shoe", 13: "face",
        14: "left-leg", 15: "right-leg", 16: "left-arm", 17: "right-arm",
        18: "bag", 19: "scarf"
    }

    def __init__(self):
        self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
        self.model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
        self.model.eval()

    def get_clothing_mask(self, image: Image.Image, clothing_type: str = "upper") -> Image.Image:
        inputs = self.processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)

        segmap = self.processor.post_process_semantic_segmentation(
            outputs, target_sizes=[image.size[::-1]]
        )[0]

        clothing_ids = {
            "upper": [4],          # верхняя одежда
            "lower": [5, 6],       # юбка, брюки
            "dress": [7],          # платье
            "full": [4, 5, 6, 7],  # вся одежда
        }

        target_ids = clothing_ids.get(clothing_type, [4])
        mask = np.zeros(segmap.shape, dtype=np.uint8)
        for label_id in target_ids:
            mask[segmap.numpy() == label_id] = 255

        return Image.fromarray(mask)

Обработка каталога одежды

class GarmentCatalogProcessor:
    """Предобработка изображений товаров для virtual try-on"""

    async def prepare_garment(self, garment_image: bytes) -> dict:
        img = Image.open(io.BytesIO(garment_image)).convert("RGB")

        # Убираем фон с одежды
        from rembg import remove
        garment_no_bg = remove(garment_image)

        # Стандартизируем размер
        img_resized = Image.open(io.BytesIO(garment_no_bg)).resize((768, 1024))

        # Генерируем описание товара через GPT-4o Vision
        description = await self.describe_garment(garment_image)

        return {
            "processed_image": img_resized,
            "description": description,
            "category": await self.classify_garment_type(garment_image)
        }

    async def describe_garment(self, garment_image: bytes) -> str:
        client = AsyncOpenAI()
        import base64
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": "Описание одежды для системы virtual try-on (материал, цвет, покрой, детали). Одно предложение, на английском."},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64.b64encode(garment_image).decode()}"}}
                ]
            }]
        )
        return response.choices[0].message.content

FastAPI сервис

from fastapi import FastAPI, File, UploadFile, Form

app = FastAPI()
tryon = VirtualTryOnService()

@app.post("/try-on")
async def virtual_try_on(
    person: UploadFile = File(...),
    garment: UploadFile = File(...),
    garment_desc: str = Form("")
):
    person_bytes = await person.read()
    garment_bytes = await garment.read()

    result = tryon.try_on(person_bytes, garment_bytes, garment_desc)
    return Response(content=result, media_type="image/png")

Метрики качества

IDM-VTON требует A100 или RTX 4090 для production-качества с разумной скоростью (~10–20 сек/изображение). RTX 3090 — ~30–40 сек. Сроки: MVP с IDM-VTON — 3–4 недели. Production-сервис с каталогом и оптимизацией — 2–3 месяца.