VTON-project

SKT FLY-AI 6기 프로젝트 Wheel Wear
AI 모델 학습 및 FastAPI 배포 Repo

wrote by Seunghyuk Choi

To-do

• docker image complete

Virtual Try-On Server

This project provides a server for a virtual try-on system, allowing users to test clothing virtually. It leverages deep learning models and is built with Python, PyTorch, and CUDA.

Prerequisites

• Python: 3.10
• CUDA: 12.4 (Ensure NVIDIA CUDA Toolkit 12.4 is installed on your system)
• Recommended: Use a Conda environment for dependency management

Setup Instructions

1. Create and Activate a Conda Environment (optional but recommended):

   conda create -n vton python=3.10
   conda activate vton

2. How to start virtual try-on FastAPI server

   pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
   git clone https://github.com/WheelWear/VTON-project.git
   pip install -r requirements.txt
   python -m uvicorn app:app --host 0.0.0.0 --port 8000
   

   • Use ngrok for secure tunnel from a public URL to your local machine

   Run this if you need original CatVTON for training

   cd VTON-project
   git clone https://github.com/Zheng-Chong/CatVTON.git
   cp pipeline_train.py CatVTON/model/
   

   Run with Docker

   docker pull coldbrew9/wheelwear-cu12.4-p3.10:latest
   docker run --gpus all -it -p 8000:8000 coldbrew9/wheelwear-cu12.4-p3.10:latest
   

llm_agent.py

• recommend_size 엔드포인트

input ex.

{
  "brand": "나이키",
  "cloth_size": "100",
  "cloth_type" : "top",
  "gender": "남성",
  "chest_circumference": 32,
  "shoulder_width": 55,
  "arm_length": 42,
  "waist_circumference": 14
}

output ex.

{
  "recommend_size": "M",
  "additional_explanation": "최근 나이키 의류가 이전보다 크게 제작된다는 리뷰가 많으며, 사용자의 일반적인 사이즈(100)와 비교했을 때 M 사이즈가 적절할 것으로 판단됩니다.",
  "references": [
    "https://www.reddit.com/r/Nike/comments/15luox2/why_does_nike_keep_making_their_shirts_bigger/",
    "https://www.trustpilot.com/review/www.nike.com",
    "https://www.today.com/shop/nike-go-firm-support-leggings-review-rcna190644"
  ],
  "reference_num": 3
}

input type, output type

class InputData(BaseModel):
    brand: str
    cloth_size: str
    cloth_type: str
    gender: str
    chest_circumference: float
    shoulder_width: float
    arm_length: float
    waist_circumference: float

class OutputData(BaseModel):
    recommend_size: str
    additional_explanation: str
    references: list
    reference_num: int

Docker

#docker pull base image
docker pull nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
#docker build
docker build -t coldbrew9/wheelwear-cu12.4-p3.10:latest .
#docker run
docker run --gpus all -it -p 8000:8000 coldbrew9/wheelwear-cu12.4-p3.10:latest

Run Train

# local train
python train_lora.py  --data_root_path ./dataset --output_dir ./experiments/ckpt --use_fp16 True --num_epochs 5 --batch_size 1 --lr 1e-4 --lora_rank 4 --accumulation_steps 4

# 코랩환경 실행
python train_lora.py  --data_root_path /content --output_dir ./experiments/ckpt --use_fp16 True --num_epochs 5 --batch_size 1 --lr 1e-4 --lora_rank 4

Prepare dataset

• unzip and rename

unzip dataset_v{version}.zip .
mv dataset_v{version} dataset

파인튜닝을 위한 데이터셋 구축 (앉아있는 자세)

dataset/
├── cloth/
│   ├── lower_img/       
│   │   └── 00000.jpg    # 하의 이미지
│   ├── lower_mask/      
│   │   └── 00000.jpg    # 하의 이미지의 마스크
│   ├── upper_img/       
│   │   └── 00000.jpg    # 하의 이미지
│   └── upper_mask/      
│       └── 00000.jpg    # 하의 이미지의 마스크
├── image/               
│   └── 00000.jpg        # 사람 이미지지
├── image_mask_L/        # 이미지의 하반신 마스크 저장 (Lower 부분)
│   └── 00000.jpg
└── image_mask_U/        # 이미지의 상반신 마스크 저장 (Upper 부분)
    └── 00000.jpg

데이터셋 구성 예시시

분류	원본 이미지 예시	마스크 이미지 예시
상의(Cloth Upper)
하의(Cloth Lower)
인물(Person)		하반신 마스크: 상반신 마스크: