model-yolo-person-classify-market/person_attr.py

"""Market-1501 Person Attribute Recognition — standalone ClearML Agent.

Deploy as two files in the same directory:
  person_attribute.py
  net_last.pth
"""

from __future__ import annotations

import argparse
import os
from io import BytesIO
from typing import Any

import requests
import torch
from clearml import Task
from PIL import Image
from torch import nn
from torch.nn import init
from torchvision import models
from torchvision import transforms as T

# ---------------------------------------------------------------------------
# Market-1501 label / attribute constants (order must match training)
# ---------------------------------------------------------------------------

NUM_LABEL = 30

MARKET_LABELS = [
    'young',
    'teenager',
    'adult',
    'old',
    'backpack',
    'bag',
    'handbag',
    'clothes',
    'down',
    'up',
    'hair',
    'hat',
    'gender',
    'upblack',
    'upwhite',
    'upred',
    'uppurple',
    'upyellow',
    'upgray',
    'upblue',
    'upgreen',
    'downblack',
    'downwhite',
    'downpink',
    'downpurple',
    'downyellow',
    'downgray',
    'downblue',
    'downgreen',
    'downbrown',
]

MARKET_ATTRIBUTES = {
    'bag': ['carrying bag', ['no', 'yes']],
    'upred': ['color of upper-body clothing', [None, 'red']],
    'upblue': ['color of upper-body clothing', [None, 'blue']],
    'hat': ['wearing hat', ['no', 'yes']],
    'downgreen': ['color of lower-body clothing', [None, 'green']],
    'downbrown': ['color of lower-body clothing', [None, 'brown']],
    'upyellow': ['color of upper-body clothing', [None, 'yellow']],
    'up': ['sleeve length', ['long sleeve', 'short sleeve']],
    'upgreen': ['color of upper-body clothing', [None, 'green']],
    'handbag': ['carrying handbag', ['no', 'yes']],
    'downgray': ['color of lower-body clothing', [None, 'gray']],
    'clothes': ['type of lower-body clothing', ['dress', 'pants']],
    'adult': ['age', [None, 'adult']],
    'downblack': ['color of lower-body clothing', [None, 'black']],
    'backpack': ['carrying backpack', ['no', 'yes']],
    'downwhite': ['color of lower-body clothing', [None, 'white']],
    'upblack': ['color of upper-body clothing', [None, 'black']],
    'gender': ['gender', ['male', 'female']],
    'downyellow': ['color of lower-body clothing', [None, 'yellow']],
    'downpink': ['color of lower-body clothing', [None, 'pink']],
    'old': ['age', [None, 'old']],
    'down': ['length of lower-body clothing', ['long lower body clothing', 'short']],
    'uppurple': ['color of upper-body clothing', [None, 'purple']],
    'downpurple': ['color of lower-body clothing', [None, 'purple']],
    'young': ['age', [None, 'young']],
    'teenager': ['age', [None, 'teenager']],
    'hair': ['hair length', ['short hair', 'long hair']],
    'downblue': ['color of lower-body clothing', [None, 'blue']],
    'upgray': ['color of upper-body clothing', [None, 'gray']],
    'upwhite': ['color of upper-body clothing', [None, 'white']],
}

TRANSFORMS = T.Compose([
    T.Resize(size=(288, 144)),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

PACKAGE_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_CHECKPOINT = os.path.join(PACKAGE_DIR, 'net_last.pth')


# ---------------------------------------------------------------------------
# Model (inference only)
# ---------------------------------------------------------------------------

def _weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
    elif classname.find('Linear') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
        init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm1d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)


def _weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        init.normal_(m.weight.data, std=0.001)
        init.constant_(m.bias.data, 0.0)


class ClassBlock(nn.Module):
    def __init__(self, input_dim, class_num=1, activ='sigmoid', num_bottleneck=512):
        super().__init__()
        add_block = nn.Sequential(
            nn.Linear(input_dim, num_bottleneck),
            nn.BatchNorm1d(num_bottleneck),
            nn.LeakyReLU(0.1),
            nn.Dropout(p=0.5),
        )
        add_block.apply(_weights_init_kaiming)

        classifier_layers = [nn.Linear(num_bottleneck, class_num)]
        if activ == 'sigmoid':
            classifier_layers.append(nn.Sigmoid())
        elif activ == 'softmax':
            classifier_layers.append(nn.Softmax())
        elif activ != 'none':
            raise AssertionError(f'Unsupported activation: {activ}')
        classifier = nn.Sequential(*classifier_layers)
        classifier.apply(_weights_init_classifier)

        self.add_block = add_block
        self.classifier = classifier

    def forward(self, x):
        x = self.add_block(x)
        return self.classifier(x)


class PersonAttributeModel(nn.Module):
    """ResNet50 + nFC, 30 Market-1501 attribute heads."""

    def __init__(self, class_num: int = NUM_LABEL):
        super().__init__()
        self.class_num = class_num
        backbone = models.resnet50(weights=None)
        backbone.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        backbone.fc = nn.Sequential()
        self.features = backbone
        self.num_ftrs = 2048
        for c in range(self.class_num):
            self.__setattr__(
                f'class_{c}',
                ClassBlock(input_dim=self.num_ftrs, class_num=1, activ='sigmoid'),
            )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        pred_label = [self.__getattr__(f'class_{c}')(x) for c in range(self.class_num)]
        return torch.cat(pred_label, dim=1)


# ---------------------------------------------------------------------------
# Processor
# ---------------------------------------------------------------------------

class PersonAttribute:
    """image_url + xywh(top-left) → Market person attributes.

    ClearML Serving: subclass with auto_load=False, then call load(model_filepath).
    """

    def __init__(
        self,
        checkpoint_path: str | None = None,
        device: str | None = None,
        auto_load: bool = True,
    ):
        self.checkpoint_path = checkpoint_path or DEFAULT_CHECKPOINT
        self.device = torch.device(
            device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        )
        self.model: PersonAttributeModel | None = None
        if auto_load:
            self.model = self.load_model()

    def load(self, model_filepath: str) -> PersonAttribute:
        """ClearML Serving hook — load weights from a downloaded checkpoint path."""
        self.checkpoint_path = model_filepath
        self.model = self.load_model()
        return self

    def load_model(self) -> PersonAttributeModel:
        if not os.path.isfile(self.checkpoint_path):
            raise FileNotFoundError(f'Checkpoint not found: {self.checkpoint_path}')
        model = PersonAttributeModel(NUM_LABEL)
        state = torch.load(self.checkpoint_path, map_location=self.device)
        model.load_state_dict(state)
        model.to(self.device)
        model.eval()
        return model

    @staticmethod
    def parse_xywh(xywh: str) -> tuple[int, int, int, int]:
        parts = str(xywh).split(',')
        if len(parts) != 4:
            raise ValueError(f"xywh must be 'x,y,w,h', got: {xywh}")
        x, y, w, h = (int(v.strip()) for v in parts)
        if w <= 0 or h <= 0:
            raise ValueError(f'xywh width/height must be positive, got: {xywh}')
        return x, y, w, h

    @staticmethod
    def load_image(image_source: str) -> Image.Image:
        if image_source.startswith(('http://', 'https://')):
            response = requests.get(image_source, timeout=60)
            response.raise_for_status()
            return Image.open(BytesIO(response.content)).convert('RGB')
        if not os.path.isfile(image_source):
            raise FileNotFoundError(f'Image not found: {image_source}')
        return Image.open(image_source).convert('RGB')

    @staticmethod
    def crop_person(image: Image.Image, xywh: tuple[int, int, int, int]) -> Image.Image:
        # Top-left based xywh (same as existing pipeline / start.py)
        x, y, w, h = xywh
        return image.crop((x, y, x + w, y + h))

    @staticmethod
    def preprocess(image: Image.Image) -> torch.Tensor:
        return TRANSFORMS(image).unsqueeze(dim=0)

    def decode(self, pred: torch.Tensor) -> dict[str, str]:
        pred = pred.squeeze(dim=0)
        results: dict[str, str] = {}
        for idx, label_key in enumerate(MARKET_LABELS):
            name, choice = MARKET_ATTRIBUTES[label_key]
            value = choice[int(pred[idx].item())]
            if value:
                results[name] = value
        return results

    def predict(self, image: Image.Image) -> dict[str, str]:
        if self.model is None:
            raise RuntimeError('Model is not loaded. Call load(model_filepath) first.')
        src = self.preprocess(image).to(self.device)
        with torch.no_grad():
            out = self.model.forward(src)
        pred = torch.gt(out, torch.ones_like(out) / 2)
        return self.decode(pred)

    def process(self, image_url: str, xywh: str) -> dict[str, str]:
        if not image_url or not str(image_url).strip():
            raise ValueError('image_url is required')
        box = self.parse_xywh(xywh)
        image = self.load_image(str(image_url).strip())
        crop = self.crop_person(image, box)
        return self.predict(crop)

    @staticmethod
    def format_result(attributes: dict[str, Any], status: str = 'PASS') -> dict[str, Any]:
        return {'output': attributes, 'status': status}

    def run(self, image_url: str, xywh: str) -> dict[str, Any]:
        attributes = self.process(image_url, xywh)
        return self.format_result(attributes)


# ---------------------------------------------------------------------------
# ClearML entry point
# ---------------------------------------------------------------------------

def upload_final_result(task: Task, result: dict) -> None:
    task.upload_artifact(name='final_result', artifact_object=result)


def main() -> None:
    # uv run python person_attr.py --image_url "https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg" --xywh "404,290,74,193"
    parser = argparse.ArgumentParser(description='Person Attribute Recognition (Market)')
    parser.add_argument('--image_url', type=str, required=True)
    parser.add_argument('--xywh', type=str, required=True)
    args = parser.parse_args()

    task = Task.init(
        project_name='Person_Attribute_Recognition',
        task_name='person_attribute',
    )
    task.connect(parser)

    engine = PersonAttribute()
    attributes = engine.process(args.image_url, args.xywh)
    result = PersonAttribute.format_result(attributes)

    print('\n' + '=' * 50)
    print('      Person Attribute Recognition 결과      ')
    print('=' * 50)
    for name, value in attributes.items():
        print(f'{name}: {value}')
    print('=' * 50)

    upload_final_result(task, result)


if __name__ == '__main__':
    main()