sync

.gitattributes

@@ -1,35 +1,34 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+mivolo/model/__pycache__/*
+mivolo/data/__pycache__/*
+mivolo/__pycache__/*
+__pycache__/*
+model/*
+rename.sh

README.md

@@ -1,13 +1,14 @@
 ---
-title: Demo
-emoji: 🌖
+title: Gender Age Detector
+emoji: 👩🧑‍🦲
 colorFrom: indigo
 colorTo: pink
 sdk: gradio
-sdk_version: 3.35.2
+sdk_version: 4.36.0
 app_file: app.py
 pinned: false
 license: apache-2.0
+arxiv: 2307.04616
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,119 @@
+import os
+import cv2
+import imghdr
+import shutil
+import warnings
+import numpy as np
+import gradio as gr
+from dataclasses import dataclass
+from mivolo.predictor import Predictor
+from utils import is_url, download_file, get_jpg_files, MODEL_DIR
+
+TMP_DIR = "./__pycache__"
+
+
+@dataclass
+class Cfg:
+    detector_weights: str
+    checkpoint: str
+    device: str = "cpu"
+    with_persons: bool = True
+    disable_faces: bool = False
+    draw: bool = True
+
+
+class ValidImgDetector:
+    predictor = None
+
+    def __init__(self):
+        detector_path = f"{MODEL_DIR}/yolov8x_person_face.pt"
+        age_gender_path = f"{MODEL_DIR}/model_imdb_cross_person_4.22_99.46.pth.tar"
+        predictor_cfg = Cfg(detector_path, age_gender_path)
+        self.predictor = Predictor(predictor_cfg)
+
+    def _detect(
+        self,
+        image: np.ndarray,
+        score_threshold: float,
+        iou_threshold: float,
+        mode: str,
+        predictor: Predictor,
+    ) -> np.ndarray:
+        # input is rgb image, output must be rgb too
+        predictor.detector.detector_kwargs["conf"] = score_threshold
+        predictor.detector.detector_kwargs["iou"] = iou_threshold
+        if mode == "Use persons and faces":
+            use_persons = True
+            disable_faces = False
+
+        elif mode == "Use persons only":
+            use_persons = True
+            disable_faces = True
+
+        elif mode == "Use faces only":
+            use_persons = False
+            disable_faces = False
+
+        predictor.age_gender_model.meta.use_persons = use_persons
+        predictor.age_gender_model.meta.disable_faces = disable_faces
+        # image = image[:, :, ::-1]  # RGB -> BGR
+        detected_objects, out_im = predictor.recognize(image)
+        has_child, has_female, has_male = False, False, False
+        if len(detected_objects.ages) > 0:
+            has_child = min(detected_objects.ages) < 18
+            has_female = "female" in detected_objects.genders
+            has_male = "male" in detected_objects.genders
+
+        return out_im[:, :, ::-1], has_child, has_female, has_male
+
+    def valid_img(self, img_path):
+        image = cv2.imread(img_path)
+        return self._detect(image, 0.4, 0.7, "Use persons and faces", self.predictor)
+
+
+def infer(photo: str):
+    if is_url(photo):
+        if os.path.exists(TMP_DIR):
+            shutil.rmtree(TMP_DIR)
+
+        photo = download_file(photo, f"{TMP_DIR}/download.jpg")
+
+    detector = ValidImgDetector()
+    if not photo or not os.path.exists(photo) or imghdr.what(photo) == None:
+        return None, None, None, "请正确输入图片 Please input image correctly"
+
+    return detector.valid_img(photo)
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as iface:
+        warnings.filterwarnings("ignore")
+        with gr.Tab("上传模式 Upload Mode"):
+            gr.Interface(
+                fn=infer,
+                inputs=gr.Image(label="上传照片 Upload Photo", type="filepath"),
+                outputs=[
+                    gr.Image(label="检测结果 Detection Result", type="numpy"),
+                    gr.Textbox(label="存在儿童 Has Child"),
+                    gr.Textbox(label="存在女性 Has Female"),
+                    gr.Textbox(label="存在男性 Has Male"),
+                ],
+                examples=get_jpg_files(f"{MODEL_DIR}/examples"),
+                allow_flagging="never",
+            )
+
+        with gr.Tab("在线模式 Online Mode"):
+            gr.Interface(
+                fn=infer,
+                inputs=gr.Textbox(label="网络图片链接 Online Picture URL"),
+                outputs=[
+                    gr.Image(label="检测结果 Detection Result", type="numpy"),
+                    gr.Textbox(label="存在儿童 Has Child"),
+                    gr.Textbox(label="存在女性 Has Female"),
+                    gr.Textbox(label="存在男性 Has Male"),
+                ],
+                allow_flagging="never",
+                cache_examples=False,
+            )
+
+    iface.launch()

mivolo/data/data_reader.py

@@ -0,0 +1,125 @@
+import os
+from collections import defaultdict
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+import pandas as pd
+
+IMAGES_EXT: Tuple = (".jpeg", ".jpg", ".png", ".webp", ".bmp", ".gif")
+VIDEO_EXT: Tuple = (".mp4", ".avi", ".mov", ".mkv", ".webm")
+
+
+@dataclass
+class PictureInfo:
+    image_path: str
+    age: Optional[str]  # age or age range(start;end format) or "-1"
+    gender: Optional[str]  # "M" of "F" or "-1"
+    bbox: List[int] = field(default_factory=lambda: [-1, -1, -1, -1])  # face bbox: xyxy
+    person_bbox: List[int] = field(default_factory=lambda: [-1, -1, -1, -1])  # person bbox: xyxy
+
+    @property
+    def has_person_bbox(self) -> bool:
+        return any(coord != -1 for coord in self.person_bbox)
+
+    @property
+    def has_face_bbox(self) -> bool:
+        return any(coord != -1 for coord in self.bbox)
+
+    def has_gt(self, only_age: bool = False) -> bool:
+        if only_age:
+            return self.age != "-1"
+        else:
+            return not (self.age == "-1" and self.gender == "-1")
+
+    def clear_person_bbox(self):
+        self.person_bbox = [-1, -1, -1, -1]
+
+    def clear_face_bbox(self):
+        self.bbox = [-1, -1, -1, -1]
+
+
+class AnnotType(Enum):
+    ORIGINAL = "original"
+    PERSONS = "persons"
+    NONE = "none"
+
+    @classmethod
+    def _missing_(cls, value):
+        print(f"WARN: Unknown annotation type {value}.")
+        return AnnotType.NONE
+
+
+def get_all_files(path: str, extensions: Tuple = IMAGES_EXT):
+    files_all = []
+    for root, subFolders, files in os.walk(path):
+        for name in files:
+            # linux tricks with .directory that still is file
+            if "directory" not in name and sum([ext.lower() in name.lower() for ext in extensions]) > 0:
+                files_all.append(os.path.join(root, name))
+    return files_all
+
+
+class InputType(Enum):
+    Image = 0
+    Video = 1
+    VideoStream = 2
+
+
+def get_input_type(input_path: str) -> InputType:
+    if os.path.isdir(input_path):
+        print("Input is a folder, only images will be processed")
+        return InputType.Image
+    elif os.path.isfile(input_path):
+        if input_path.endswith(VIDEO_EXT):
+            return InputType.Video
+        if input_path.endswith(IMAGES_EXT):
+            return InputType.Image
+        else:
+            raise ValueError(
+                f"Unknown or unsupported input file format {input_path}, \
+                             supported video formats: {VIDEO_EXT}, \
+                             supported image formats: {IMAGES_EXT}"
+            )
+    elif input_path.startswith("http") and not input_path.endswith(IMAGES_EXT):
+        return InputType.VideoStream
+    else:
+        raise ValueError(f"Unknown input {input_path}")
+
+
+def read_csv_annotation_file(annotation_file: str, images_dir: str, ignore_without_gt=False):
+    bboxes_per_image: Dict[str, List[PictureInfo]] = defaultdict(list)
+
+    df = pd.read_csv(annotation_file, sep=",")
+
+    annot_type = AnnotType("persons") if "person_x0" in df.columns else AnnotType("original")
+    print(f"Reading {annotation_file} (type: {annot_type})...")
+
+    missing_images = 0
+    for index, row in df.iterrows():
+        img_path = os.path.join(images_dir, row["img_name"])
+        if not os.path.exists(img_path):
+            missing_images += 1
+            continue
+
+        face_x1, face_y1, face_x2, face_y2 = row["face_x0"], row["face_y0"], row["face_x1"], row["face_y1"]
+        age, gender = str(row["age"]), str(row["gender"])
+
+        if ignore_without_gt and (age == "-1" or gender == "-1"):
+            continue
+
+        if annot_type == AnnotType.PERSONS:
+            p_x1, p_y1, p_x2, p_y2 = row["person_x0"], row["person_y0"], row["person_x1"], row["person_y1"]
+            person_bbox = list(map(int, [p_x1, p_y1, p_x2, p_y2]))
+        else:
+            person_bbox = [-1, -1, -1, -1]
+
+        bbox = list(map(int, [face_x1, face_y1, face_x2, face_y2]))
+        pic_info = PictureInfo(img_path, age, gender, bbox, person_bbox)
+        assert isinstance(pic_info.person_bbox, list)
+
+        bboxes_per_image[img_path].append(pic_info)
+
+    if missing_images > 0:
+        print(f"WARNING: Missing images: {missing_images}/{len(df)}")
+    return bboxes_per_image, annot_type

mivolo/data/dataset/__init__.py

@@ -0,0 +1,64 @@
+from typing import Tuple
+
+import torch
+from mivolo.model.mi_volo import MiVOLO
+
+from .age_gender_dataset import AgeGenderDataset
+from .age_gender_loader import create_loader
+from .classification_dataset import AdienceDataset, FairFaceDataset
+
+DATASET_CLASS_MAP = {
+    "utk": AgeGenderDataset,
+    "lagenda": AgeGenderDataset,
+    "imdb": AgeGenderDataset,
+    "adience": AdienceDataset,
+    "fairface": FairFaceDataset,
+}
+
+
+def build(
+    name: str,
+    images_path: str,
+    annotations_path: str,
+    split: str,
+    mivolo_model: MiVOLO,
+    workers: int,
+    batch_size: int,
+) -> Tuple[torch.utils.data.Dataset, torch.utils.data.DataLoader]:
+
+    dataset_class = DATASET_CLASS_MAP[name]
+
+    dataset: torch.utils.data.Dataset = dataset_class(
+        images_path=images_path,
+        annotations_path=annotations_path,
+        name=name,
+        split=split,
+        target_size=mivolo_model.input_size,
+        max_age=mivolo_model.meta.max_age,
+        min_age=mivolo_model.meta.min_age,
+        model_with_persons=mivolo_model.meta.with_persons_model,
+        use_persons=mivolo_model.meta.use_persons,
+        disable_faces=mivolo_model.meta.disable_faces,
+        only_age=mivolo_model.meta.only_age,
+    )
+
+    data_config = mivolo_model.data_config
+
+    in_chans = 3 if not mivolo_model.meta.with_persons_model else 6
+    input_size = (in_chans, mivolo_model.input_size, mivolo_model.input_size)
+
+    dataset_loader: torch.utils.data.DataLoader = create_loader(
+        dataset,
+        input_size=input_size,
+        batch_size=batch_size,
+        mean=data_config["mean"],
+        std=data_config["std"],
+        num_workers=workers,
+        crop_pct=data_config["crop_pct"],
+        crop_mode=data_config["crop_mode"],
+        pin_memory=False,
+        device=mivolo_model.device,
+        target_type=dataset.target_dtype,
+    )
+
+    return dataset, dataset_loader

mivolo/data/dataset/age_gender_dataset.py

@@ -0,0 +1,194 @@
+import logging
+from typing import Any, List, Optional, Set
+
+import cv2
+import numpy as np
+import torch
+from mivolo.data.dataset.reader_age_gender import ReaderAgeGender
+from PIL import Image
+from torchvision import transforms
+
+_logger = logging.getLogger("AgeGenderDataset")
+
+
+class AgeGenderDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        images_path,
+        annotations_path,
+        name=None,
+        split="train",
+        load_bytes=False,
+        img_mode="RGB",
+        transform=None,
+        is_training=False,
+        seed=1234,
+        target_size=224,
+        min_age=None,
+        max_age=None,
+        model_with_persons=False,
+        use_persons=False,
+        disable_faces=False,
+        only_age=False,
+    ):
+        reader = ReaderAgeGender(
+            images_path,
+            annotations_path,
+            split=split,
+            seed=seed,
+            target_size=target_size,
+            with_persons=use_persons,
+            disable_faces=disable_faces,
+            only_age=only_age,
+        )
+
+        self.name = name
+        self.model_with_persons = model_with_persons
+        self.reader = reader
+        self.load_bytes = load_bytes
+        self.img_mode = img_mode
+        self.transform = transform
+        self._consecutive_errors = 0
+        self.is_training = is_training
+        self.random_flip = 0.0
+
+        # Setting up classes.
+        # If min and max classes are passed - use them to have the same preprocessing for validation
+        self.max_age: float = None
+        self.min_age: float = None
+        self.avg_age: float = None
+        self.set_ages_min_max(min_age, max_age)
+
+        self.genders = ["M", "F"]
+        self.num_classes_gender = len(self.genders)
+
+        self.age_classes: Optional[List[str]] = self.set_age_classes()
+
+        self.num_classes_age = 1 if self.age_classes is None else len(self.age_classes)
+        self.num_classes: int = self.num_classes_age + self.num_classes_gender
+        self.target_dtype = torch.float32
+
+    def set_age_classes(self) -> Optional[List[str]]:
+        return None  # for regression dataset
+
+    def set_ages_min_max(self, min_age: Optional[float], max_age: Optional[float]):
+
+        assert all(age is None for age in [min_age, max_age]) or all(
+            age is not None for age in [min_age, max_age]
+        ), "Both min and max age must be passed or none of them"
+
+        if max_age is not None and min_age is not None:
+            _logger.info(f"Received predefined min_age {min_age} and max_age {max_age}")
+            self.max_age = max_age
+            self.min_age = min_age
+        else:
+            # collect statistics from loaded dataset
+            all_ages_set: Set[int] = set()
+            for img_path, image_samples in self.reader._ann.items():
+                for image_sample_info in image_samples:
+                    if image_sample_info.age == "-1":
+                        continue
+                    age = round(float(image_sample_info.age))
+                    all_ages_set.add(age)
+
+            self.max_age = max(all_ages_set)
+            self.min_age = min(all_ages_set)
+
+        self.avg_age = (self.max_age + self.min_age) / 2.0
+
+    def _norm_age(self, age):
+        return (age - self.avg_age) / (self.max_age - self.min_age)
+
+    def parse_gender(self, _gender: str) -> float:
+        if _gender != "-1":
+            gender = float(0 if _gender == "M" or _gender == "0" else 1)
+        else:
+            gender = -1
+        return gender
+
+    def parse_target(self, _age: str, gender: str) -> List[Any]:
+        if _age != "-1":
+            age = round(float(_age))
+            age = self._norm_age(float(age))
+        else:
+            age = -1
+
+        target: List[float] = [age, self.parse_gender(gender)]
+        return target
+
+    @property
+    def transform(self):
+        return self._transform
+
+    @transform.setter
+    def transform(self, transform):
+        # Disable pretrained monkey-patched transforms
+        if not transform:
+            return
+
+        _trans = []
+        for trans in transform.transforms:
+            if "Resize" in str(trans):
+                continue
+            if "Crop" in str(trans):
+                continue
+            _trans.append(trans)
+        self._transform = transforms.Compose(_trans)
+
+    def apply_tranforms(self, image: Optional[np.ndarray]) -> np.ndarray:
+        if image is None:
+            return None
+
+        if self.transform is None:
+            return image
+
+        image = convert_to_pil(image, self.img_mode)
+        for trans in self.transform.transforms:
+            image = trans(image)
+        return image
+
+    def __getitem__(self, index):
+        # get preprocessed face and person crops (np.ndarray)
+        # resize + pad, for person crops: cut off other bboxes
+        images, target = self.reader[index]
+
+        target = self.parse_target(*target)
+
+        if self.model_with_persons:
+            face_image, person_image = images
+            person_image: np.ndarray = self.apply_tranforms(person_image)
+        else:
+            face_image = images[0]
+            person_image = None
+
+        face_image: np.ndarray = self.apply_tranforms(face_image)
+
+        if person_image is not None:
+            img = np.concatenate([face_image, person_image], axis=0)
+        else:
+            img = face_image
+
+        return img, target
+
+    def __len__(self):
+        return len(self.reader)
+
+    def filename(self, index, basename=False, absolute=False):
+        return self.reader.filename(index, basename, absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return self.reader.filenames(basename, absolute)
+
+
+def convert_to_pil(cv_im: Optional[np.ndarray], img_mode: str = "RGB") -> "Image":
+    if cv_im is None:
+        return None
+
+    if img_mode == "RGB":
+        cv_im = cv2.cvtColor(cv_im, cv2.COLOR_BGR2RGB)
+    else:
+        raise Exception("Incorrect image mode has been passed!")
+
+    cv_im = np.ascontiguousarray(cv_im)
+    pil_image = Image.fromarray(cv_im)
+    return pil_image

mivolo/data/dataset/age_gender_loader.py

@@ -0,0 +1,169 @@
+"""
+Code adapted from timm https://github.com/huggingface/pytorch-image-models
+
+Modifications and additions for mivolo by / Copyright 2023, Irina Tolstykh, Maxim Kuprashevich
+"""
+
+import logging
+from contextlib import suppress
+from functools import partial
+from itertools import repeat
+
+import numpy as np
+import torch
+import torch.utils.data
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.data.dataset import IterableImageDataset
+from timm.data.loader import PrefetchLoader, _worker_init
+from timm.data.transforms_factory import create_transform
+
+_logger = logging.getLogger(__name__)
+
+
+def fast_collate(batch, target_dtype=torch.uint8):
+    """A fast collation function optimized for uint8 images (np array or torch) and target_dtype targets (labels)"""
+    assert isinstance(batch[0], tuple)
+    batch_size = len(batch)
+    if isinstance(batch[0][0], np.ndarray):
+        targets = torch.tensor([b[1] for b in batch], dtype=target_dtype)
+        assert len(targets) == batch_size
+        tensor = torch.zeros((batch_size, *batch[0][0].shape), dtype=torch.uint8)
+        for i in range(batch_size):
+            tensor[i] += torch.from_numpy(batch[i][0])
+        return tensor, targets
+    else:
+        raise ValueError(f"Incorrect batch type: {type(batch[0][0])}")
+
+
+def adapt_to_chs(x, n):
+    if not isinstance(x, (tuple, list)):
+        x = tuple(repeat(x, n))
+    elif len(x) != n:
+        # doubled channels
+        if len(x) * 2 == n:
+            x = np.concatenate((x, x))
+            _logger.warning(f"Pretrained mean/std different shape than model (doubled channes), using concat: {x}.")
+        else:
+            x_mean = np.mean(x).item()
+            x = (x_mean,) * n
+            _logger.warning(f"Pretrained mean/std different shape than model, using avg value {x}.")
+    else:
+        assert len(x) == n, "normalization stats must match image channels"
+    return x
+
+
+class PrefetchLoaderForMultiInput(PrefetchLoader):
+    def __init__(
+        self,
+        loader,
+        mean=IMAGENET_DEFAULT_MEAN,
+        std=IMAGENET_DEFAULT_STD,
+        channels=3,
+        device=torch.device("cpu"),
+        img_dtype=torch.float32,
+    ):
+
+        mean = adapt_to_chs(mean, channels)
+        std = adapt_to_chs(std, channels)
+        normalization_shape = (1, channels, 1, 1)
+
+        self.loader = loader
+        self.device = device
+        self.img_dtype = img_dtype
+        self.mean = torch.tensor([x * 255 for x in mean], device=device, dtype=img_dtype).view(normalization_shape)
+        self.std = torch.tensor([x * 255 for x in std], device=device, dtype=img_dtype).view(normalization_shape)
+
+        self.is_cuda = torch.cuda.is_available() and device.type == "cpu"
+
+    def __iter__(self):
+        first = True
+        if self.is_cuda:
+            stream = torch.cuda.Stream()
+            stream_context = partial(torch.cuda.stream, stream=stream)
+        else:
+            stream = None
+            stream_context = suppress
+
+        for next_input, next_target in self.loader:
+
+            with stream_context():
+                next_input = next_input.to(device=self.device, non_blocking=True)
+                next_target = next_target.to(device=self.device, non_blocking=True)
+                next_input = next_input.to(self.img_dtype).sub_(self.mean).div_(self.std)
+
+            if not first:
+                yield input, target  # noqa: F823, F821
+            else:
+                first = False
+
+            if stream is not None:
+                torch.cuda.current_stream().wait_stream(stream)
+
+            input = next_input
+            target = next_target
+
+        yield input, target
+
+
+def create_loader(
+    dataset,
+    input_size,
+    batch_size,
+    mean=IMAGENET_DEFAULT_MEAN,
+    std=IMAGENET_DEFAULT_STD,
+    num_workers=1,
+    crop_pct=None,
+    crop_mode=None,
+    pin_memory=False,
+    img_dtype=torch.float32,
+    device=torch.device("cpu"),
+    persistent_workers=True,
+    worker_seeding="all",
+    target_type=torch.int64,
+):
+
+    transform = create_transform(
+        input_size,
+        is_training=False,
+        use_prefetcher=True,
+        mean=mean,
+        std=std,
+        crop_pct=crop_pct,
+        crop_mode=crop_mode,
+    )
+    dataset.transform = transform
+
+    if isinstance(dataset, IterableImageDataset):
+        # give Iterable datasets early knowledge of num_workers so that sample estimates
+        # are correct before worker processes are launched
+        dataset.set_loader_cfg(num_workers=num_workers)
+        raise ValueError("Incorrect dataset type: IterableImageDataset")
+
+    loader_class = torch.utils.data.DataLoader
+    loader_args = dict(
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        sampler=None,
+        collate_fn=lambda batch: fast_collate(batch, target_dtype=target_type),
+        pin_memory=pin_memory,
+        drop_last=False,
+        worker_init_fn=partial(_worker_init, worker_seeding=worker_seeding),
+        persistent_workers=persistent_workers,
+    )
+    try:
+        loader = loader_class(dataset, **loader_args)
+    except TypeError:
+        loader_args.pop("persistent_workers")  # only in Pytorch 1.7+
+        loader = loader_class(dataset, **loader_args)
+
+    loader = PrefetchLoaderForMultiInput(
+        loader,
+        mean=mean,
+        std=std,
+        channels=input_size[0],
+        device=device,
+        img_dtype=img_dtype,
+    )
+
+    return loader

mivolo/data/dataset/classification_dataset.py

@@ -0,0 +1,48 @@
+from typing import Any, List, Optional
+
+import torch
+
+from .age_gender_dataset import AgeGenderDataset
+
+
+class ClassificationDataset(AgeGenderDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.target_dtype = torch.int32
+
+    def set_age_classes(self) -> Optional[List[str]]:
+        raise NotImplementedError
+
+    def parse_target(self, age: str, gender: str) -> List[Any]:
+        assert self.age_classes is not None
+        if age != "-1":
+            assert age in self.age_classes, f"Unknown category in {self.name} dataset: {age}"
+            age_ind = self.age_classes.index(age)
+        else:
+            age_ind = -1
+
+        target: List[int] = [age_ind, int(self.parse_gender(gender))]
+        return target
+
+
+class FairFaceDataset(ClassificationDataset):
+    def set_age_classes(self) -> Optional[List[str]]:
+        age_classes = ["0;2", "3;9", "10;19", "20;29", "30;39", "40;49", "50;59", "60;69", "70;120"]
+        # a[i-1] <= v < a[i] => age_classes[i-1]
+        self._intervals = torch.tensor([0, 3, 10, 20, 30, 40, 50, 60, 70])
+
+        return age_classes
+
+
+class AdienceDataset(ClassificationDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.target_dtype = torch.int32
+
+    def set_age_classes(self) -> Optional[List[str]]:
+        age_classes = ["0;2", "4;6", "8;12", "15;20", "25;32", "38;43", "48;53", "60;100"]
+        # a[i-1] <= v < a[i] => age_classes[i-1]
+        self._intervals = torch.tensor([0, 4, 7, 14, 24, 36, 46, 57])
+        return age_classes

mivolo/data/dataset/reader_age_gender.py

@@ -0,0 +1,490 @@
+import logging
+import os
+from functools import partial
+from multiprocessing.pool import ThreadPool
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+from mivolo.data.data_reader import AnnotType, PictureInfo, get_all_files, read_csv_annotation_file
+from mivolo.data.misc import IOU, class_letterbox, cropout_black_parts
+from timm.data.readers.reader import Reader
+from tqdm import tqdm
+
+CROP_ROUND_TOL = 0.3
+MIN_PERSON_SIZE = 100
+MIN_PERSON_CROP_AFTERCUT_RATIO = 0.4
+
+_logger = logging.getLogger("ReaderAgeGender")
+
+
+class ReaderAgeGender(Reader):
+    """
+    Reader for almost original imdb-wiki cleaned dataset.
+    Two changes:
+        1. Your annotation must be in ./annotation subdir of dataset root
+        2. Images must be in images subdir
+
+    """
+
+    def __init__(
+        self,
+        images_path,
+        annotations_path,
+        split="validation",
+        target_size=224,
+        min_size=5,
+        seed=1234,
+        with_persons=False,
+        min_person_size=MIN_PERSON_SIZE,
+        disable_faces=False,
+        only_age=False,
+        min_person_aftercut_ratio=MIN_PERSON_CROP_AFTERCUT_RATIO,
+        crop_round_tol=CROP_ROUND_TOL,
+    ):
+        super().__init__()
+
+        self.with_persons = with_persons
+        self.disable_faces = disable_faces
+        self.only_age = only_age
+
+        # can be only black for now, even though it's not very good with further normalization
+        self.crop_out_color = (0, 0, 0)
+
+        self.empty_crop = np.ones((target_size, target_size, 3)) * self.crop_out_color
+        self.empty_crop = self.empty_crop.astype(np.uint8)
+
+        self.min_person_size = min_person_size
+        self.min_person_aftercut_ratio = min_person_aftercut_ratio
+        self.crop_round_tol = crop_round_tol
+
+        self.split = split
+        self.min_size = min_size
+        self.seed = seed
+        self.target_size = target_size
+
+        # Reading annotations. Can be multiple files if annotations_path dir
+        self._ann: Dict[str, List[PictureInfo]] = {}  # list of samples for each image
+        self._associated_objects: Dict[str, Dict[int, List[List[int]]]] = {}
+        self._faces_list: List[Tuple[str, int]] = []  # samples from this list will be loaded in __getitem__
+
+        self._read_annotations(images_path, annotations_path)
+        _logger.info(f"Dataset length: {len(self._faces_list)} crops")
+
+    def __getitem__(self, index):
+        return self._read_img_and_label(index)
+
+    def __len__(self):
+        return len(self._faces_list)
+
+    def _filename(self, index, basename=False, absolute=False):
+        img_p = self._faces_list[index][0]
+        return os.path.basename(img_p) if basename else img_p
+
+    def _read_annotations(self, images_path, csvs_path):
+        self._ann = {}
+        self._faces_list = []
+        self._associated_objects = {}
+
+        csvs = get_all_files(csvs_path, [".csv"])
+        csvs = [c for c in csvs if self.split in os.path.basename(c)]
+
+        # load annotations per image
+        for csv in csvs:
+            db, ann_type = read_csv_annotation_file(csv, images_path)
+            if self.with_persons and ann_type != AnnotType.PERSONS:
+                raise ValueError(
+                    f"Annotation type in file {csv} contains no persons, "
+                    f"but annotations with persons are requested."
+                )
+            self._ann.update(db)
+
+        if len(self._ann) == 0:
+            raise ValueError("Annotations are empty!")
+
+        self._ann, self._associated_objects = self.prepare_annotations()
+        images_list = list(self._ann.keys())
+
+        for img_path in images_list:
+            for index, image_sample_info in enumerate(self._ann[img_path]):
+                assert image_sample_info.has_gt(
+                    self.only_age
+                ), "Annotations must be checked with self.prepare_annotations() func"
+                self._faces_list.append((img_path, index))
+
+    def _read_img_and_label(self, index):
+        if not isinstance(index, int):
+            raise TypeError("ReaderAgeGender expected index to be integer")
+
+        img_p, face_index = self._faces_list[index]
+        ann: PictureInfo = self._ann[img_p][face_index]
+        img = cv2.imread(img_p)
+
+        face_empty = True
+        if ann.has_face_bbox and not (self.with_persons and self.disable_faces):
+            face_crop, face_empty = self._get_crop(ann.bbox, img)
+
+        if not self.with_persons and face_empty:
+            # model without persons
+            raise ValueError("Annotations must be checked with self.prepare_annotations() func")
+
+        if face_empty:
+            face_crop = self.empty_crop
+
+        person_empty = True
+        if self.with_persons or self.disable_faces:
+            if ann.has_person_bbox:
+                # cut off all associated objects from person crop
+                objects = self._associated_objects[img_p][face_index]
+                person_crop, person_empty = self._get_crop(
+                    ann.person_bbox,
+                    img,
+                    crop_out_color=self.crop_out_color,
+                    asced_objects=objects,
+                )
+
+            if face_empty and person_empty:
+                raise ValueError("Annotations must be checked with self.prepare_annotations() func")
+
+        if person_empty:
+            person_crop = self.empty_crop
+
+        return (face_crop, person_crop), [ann.age, ann.gender]
+
+    def _get_crop(
+        self,
+        bbox,
+        img,
+        asced_objects=None,
+        crop_out_color=(0, 0, 0),
+    ) -> Tuple[np.ndarray, bool]:
+
+        empty_bbox = False
+
+        xmin, ymin, xmax, ymax = bbox
+        assert not (
+            ymax - ymin < self.min_size or xmax - xmin < self.min_size
+        ), "Annotations must be checked with self.prepare_annotations() func"
+
+        crop = img[ymin:ymax, xmin:xmax]
+
+        if asced_objects:
+            # cut off other objects for person crop
+            crop, empty_bbox = _cropout_asced_objs(
+                asced_objects,
+                bbox,
+                crop.copy(),
+                crop_out_color=crop_out_color,
+                min_person_size=self.min_person_size,
+                crop_round_tol=self.crop_round_tol,
+                min_person_aftercut_ratio=self.min_person_aftercut_ratio,
+            )
+            if empty_bbox:
+                crop = self.empty_crop
+
+        crop = class_letterbox(crop, new_shape=(self.target_size, self.target_size), color=crop_out_color)
+        return crop, empty_bbox
+
+    def prepare_annotations(self):
+
+        good_anns: Dict[str, List[PictureInfo]] = {}
+        all_associated_objects: Dict[str, Dict[int, List[List[int]]]] = {}
+
+        if not self.with_persons:
+            # remove all persons
+            for img_path, bboxes in self._ann.items():
+                for sample in bboxes:
+                    sample.clear_person_bbox()
+
+        # check dataset and collect associated_objects
+        verify_images_func = partial(
+            verify_images,
+            min_size=self.min_size,
+            min_person_size=self.min_person_size,
+            with_persons=self.with_persons,
+            disable_faces=self.disable_faces,
+            crop_round_tol=self.crop_round_tol,
+            min_person_aftercut_ratio=self.min_person_aftercut_ratio,
+            only_age=self.only_age,
+        )
+        num_threads = min(8, os.cpu_count())
+
+        all_msgs = []
+        broken = 0
+        skipped = 0
+        all_skipped_crops = 0
+        desc = "Check annotations..."
+        with ThreadPool(num_threads) as pool:
+            pbar = tqdm(
+                pool.imap_unordered(verify_images_func, list(self._ann.items())),
+                desc=desc,
+                total=len(self._ann),
+            )
+
+            for (img_info, associated_objects, msgs, is_corrupted, is_empty_annotations, skipped_crops) in pbar:
+                broken += 1 if is_corrupted else 0
+                all_msgs.extend(msgs)
+                all_skipped_crops += skipped_crops
+                skipped += 1 if is_empty_annotations else 0
+                if img_info is not None:
+                    img_path, img_samples = img_info
+                    good_anns[img_path] = img_samples
+                    all_associated_objects.update({img_path: associated_objects})
+
+                pbar.desc = (
+                    f"{desc} {skipped} images skipped ({all_skipped_crops} crops are incorrect); "
+                    f"{broken} images corrupted"
+                )
+
+            pbar.close()
+
+        for msg in all_msgs:
+            print(msg)
+        print(f"\nLeft images: {len(good_anns)}")
+
+        return good_anns, all_associated_objects
+
+
+def verify_images(
+    img_info,
+    min_size: int,
+    min_person_size: int,
+    with_persons: bool,
+    disable_faces: bool,
+    crop_round_tol: float,
+    min_person_aftercut_ratio: float,
+    only_age: bool,
+):
+    # If crop is too small, if image can not be read or if image does not exist
+    # then filter out this sample
+
+    disable_faces = disable_faces and with_persons
+    kwargs = dict(
+        min_person_size=min_person_size,
+        disable_faces=disable_faces,
+        with_persons=with_persons,
+        crop_round_tol=crop_round_tol,
+        min_person_aftercut_ratio=min_person_aftercut_ratio,
+        only_age=only_age,
+    )
+
+    def bbox_correct(bbox, min_size, im_h, im_w) -> Tuple[bool, List[int]]:
+        ymin, ymax, xmin, xmax = _correct_bbox(bbox, im_h, im_w)
+        crop_h, crop_w = ymax - ymin, xmax - xmin
+        if crop_h < min_size or crop_w < min_size:
+            return False, [-1, -1, -1, -1]
+        bbox = [xmin, ymin, xmax, ymax]
+        return True, bbox
+
+    msgs = []
+    skipped_crops = 0
+    is_corrupted = False
+    is_empty_annotations = False
+
+    img_path: str = img_info[0]
+    img_samples: List[PictureInfo] = img_info[1]
+    try:
+        im_cv = cv2.imread(img_path)
+        im_h, im_w = im_cv.shape[:2]
+    except Exception:
+        msgs.append(f"Can not load image {img_path}")
+        is_corrupted = True
+        return None, {}, msgs, is_corrupted, is_empty_annotations, skipped_crops
+
+    out_samples: List[PictureInfo] = []
+    for sample in img_samples:
+        # correct face bbox
+        if sample.has_face_bbox:
+            is_correct, sample.bbox = bbox_correct(sample.bbox, min_size, im_h, im_w)
+            if not is_correct and sample.has_gt(only_age):
+                msgs.append("Small face. Passing..")
+                skipped_crops += 1
+
+        # correct person bbox
+        if sample.has_person_bbox:
+            is_correct, sample.person_bbox = bbox_correct(
+                sample.person_bbox, max(min_person_size, min_size), im_h, im_w
+            )
+            if not is_correct and sample.has_gt(only_age):
+                msgs.append(f"Small person {img_path}. Passing..")
+                skipped_crops += 1
+
+        if sample.has_face_bbox or sample.has_person_bbox:
+            out_samples.append(sample)
+        elif sample.has_gt(only_age):
+            msgs.append("Sample hs no face and no body. Passing..")
+            skipped_crops += 1
+
+    # sort that samples with undefined age and gender be the last
+    out_samples = sorted(out_samples, key=lambda sample: 1 if not sample.has_gt(only_age) else 0)
+
+    # for each person find other faces and persons bboxes, intersected with it
+    associated_objects: Dict[int, List[List[int]]] = find_associated_objects(out_samples, only_age=only_age)
+
+    out_samples, associated_objects, skipped_crops = filter_bad_samples(
+        out_samples, associated_objects, im_cv, msgs, skipped_crops, **kwargs
+    )
+
+    out_img_info: Optional[Tuple[str, List]] = (img_path, out_samples)
+    if len(out_samples) == 0:
+        out_img_info = None
+        is_empty_annotations = True
+
+    return out_img_info, associated_objects, msgs, is_corrupted, is_empty_annotations, skipped_crops
+
+
+def filter_bad_samples(
+    out_samples: List[PictureInfo],
+    associated_objects: dict,
+    im_cv: np.ndarray,
+    msgs: List[str],
+    skipped_crops: int,
+    **kwargs,
+):
+    with_persons, disable_faces, min_person_size, crop_round_tol, min_person_aftercut_ratio, only_age = (
+        kwargs["with_persons"],
+        kwargs["disable_faces"],
+        kwargs["min_person_size"],
+        kwargs["crop_round_tol"],
+        kwargs["min_person_aftercut_ratio"],
+        kwargs["only_age"],
+    )
+
+    # left only samples with annotations
+    inds = [sample_ind for sample_ind, sample in enumerate(out_samples) if sample.has_gt(only_age)]
+    out_samples, associated_objects = _filter_by_ind(out_samples, associated_objects, inds)
+
+    if kwargs["disable_faces"]:
+        # clear all faces
+        for ind, sample in enumerate(out_samples):
+            sample.clear_face_bbox()
+
+        # left only samples with person_bbox
+        inds = [sample_ind for sample_ind, sample in enumerate(out_samples) if sample.has_person_bbox]
+        out_samples, associated_objects = _filter_by_ind(out_samples, associated_objects, inds)
+
+    if with_persons or disable_faces:
+        # check that preprocessing func
+        # _cropout_asced_objs() return not empty person_image for each out sample
+
+        inds = []
+        for ind, sample in enumerate(out_samples):
+            person_empty = True
+            if sample.has_person_bbox:
+                xmin, ymin, xmax, ymax = sample.person_bbox
+                crop = im_cv[ymin:ymax, xmin:xmax]
+                # cut off all associated objects from person crop
+                _, person_empty = _cropout_asced_objs(
+                    associated_objects[ind],
+                    sample.person_bbox,
+                    crop.copy(),
+                    min_person_size=min_person_size,
+                    crop_round_tol=crop_round_tol,
+                    min_person_aftercut_ratio=min_person_aftercut_ratio,
+                )
+
+            if person_empty and not sample.has_face_bbox:
+                msgs.append("Small person after preprocessing. Passing..")
+                skipped_crops += 1
+            else:
+                inds.append(ind)
+        out_samples, associated_objects = _filter_by_ind(out_samples, associated_objects, inds)
+
+    assert len(associated_objects) == len(out_samples)
+    return out_samples, associated_objects, skipped_crops
+
+
+def _filter_by_ind(out_samples, associated_objects, inds):
+    _associated_objects = {}
+    _out_samples = []
+    for ind, sample in enumerate(out_samples):
+        if ind in inds:
+            _associated_objects[len(_out_samples)] = associated_objects[ind]
+            _out_samples.append(sample)
+
+    return _out_samples, _associated_objects
+
+
+def find_associated_objects(
+    image_samples: List[PictureInfo], iou_thresh=0.0001, only_age=False
+) -> Dict[int, List[List[int]]]:
+    """
+    For each person (which has gt age and gt gender) find other faces and persons bboxes, intersected with it
+    """
+    associated_objects: Dict[int, List[List[int]]] = {}
+
+    for iindex, image_sample_info in enumerate(image_samples):
+        # add own face
+        associated_objects[iindex] = [image_sample_info.bbox] if image_sample_info.has_face_bbox else []
+
+        if not image_sample_info.has_person_bbox or not image_sample_info.has_gt(only_age):
+            # if sample has not gt => not be used
+            continue
+
+        iperson_box = image_sample_info.person_bbox
+        for jindex, other_image_sample in enumerate(image_samples):
+            if iindex == jindex:
+                continue
+            if other_image_sample.has_face_bbox:
+                jface_bbox = other_image_sample.bbox
+                iou = _get_iou(jface_bbox, iperson_box)
+                if iou >= iou_thresh:
+                    associated_objects[iindex].append(jface_bbox)
+            if other_image_sample.has_person_bbox:
+                jperson_bbox = other_image_sample.person_bbox
+                iou = _get_iou(jperson_bbox, iperson_box)
+                if iou >= iou_thresh:
+                    associated_objects[iindex].append(jperson_bbox)
+
+    return associated_objects
+
+
+def _cropout_asced_objs(
+    asced_objects,
+    person_bbox,
+    crop,
+    min_person_size,
+    crop_round_tol,
+    min_person_aftercut_ratio,
+    crop_out_color=(0, 0, 0),
+):
+    empty = False
+    xmin, ymin, xmax, ymax = person_bbox
+
+    for a_obj in asced_objects:
+        aobj_xmin, aobj_ymin, aobj_xmax, aobj_ymax = a_obj
+
+        aobj_ymin = int(max(aobj_ymin - ymin, 0))
+        aobj_xmin = int(max(aobj_xmin - xmin, 0))
+        aobj_ymax = int(min(aobj_ymax - ymin, ymax - ymin))
+        aobj_xmax = int(min(aobj_xmax - xmin, xmax - xmin))
+
+        crop[aobj_ymin:aobj_ymax, aobj_xmin:aobj_xmax] = crop_out_color
+
+    crop, cropped_ratio = cropout_black_parts(crop, crop_round_tol)
+    if (
+        crop.shape[0] < min_person_size or crop.shape[1] < min_person_size
+    ) or cropped_ratio < min_person_aftercut_ratio:
+        crop = None
+        empty = True
+
+    return crop, empty
+
+
+def _correct_bbox(bbox, h, w):
+    xmin, ymin, xmax, ymax = bbox
+    ymin = min(max(ymin, 0), h)
+    ymax = min(max(ymax, 0), h)
+    xmin = min(max(xmin, 0), w)
+    xmax = min(max(xmax, 0), w)
+    return ymin, ymax, xmin, xmax
+
+
+def _get_iou(bbox1, bbox2):
+    xmin1, ymin1, xmax1, ymax1 = bbox1
+    xmin2, ymin2, xmax2, ymax2 = bbox2
+    iou = IOU(
+        [ymin1, xmin1, ymax1, xmax1],
+        [ymin2, xmin2, ymax2, xmax2],
+    )
+    return iou

mivolo/data/misc.py

@@ -0,0 +1,264 @@
+import argparse
+import ast
+import re
+from typing import List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from scipy.optimize import linear_sum_assignment
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+CROP_ROUND_RATE = 0.1
+MIN_PERSON_CROP_NONZERO = 0.5
+
+
+def aggregate_votes_winsorized(ages, max_age_dist=6):
+    # Replace any annotation that is more than a max_age_dist away from the median
+    # with the median + max_age_dist if higher or max_age_dist - max_age_dist if below
+    median = np.median(ages)
+    ages = np.clip(ages, median - max_age_dist, median + max_age_dist)
+    return np.mean(ages)
+
+
+def cropout_black_parts(img, tol=0.3):
+    # Create a binary mask of zero pixels
+    zero_pixels_mask = np.all(img == 0, axis=2)
+    # Calculate the threshold for zero pixels in rows and columns
+    threshold = img.shape[0] - img.shape[0] * tol
+    # Calculate row sums and column sums of zero pixels mask
+    row_sums = np.sum(zero_pixels_mask, axis=1)
+    col_sums = np.sum(zero_pixels_mask, axis=0)
+    # Find the first and last rows with zero pixel sums above the threshold
+    start_row = np.argmin(row_sums > threshold)
+    end_row = img.shape[0] - np.argmin(row_sums[::-1] > threshold)
+    # Find the first and last columns with zero pixel sums above the threshold
+    start_col = np.argmin(col_sums > threshold)
+    end_col = img.shape[1] - np.argmin(col_sums[::-1] > threshold)
+    # Crop the image
+    cropped_img = img[start_row:end_row, start_col:end_col, :]
+    area = cropped_img.shape[0] * cropped_img.shape[1]
+    area_orig = img.shape[0] * img.shape[1]
+    return cropped_img, area / area_orig
+
+
+def natural_key(string_):
+    """See http://www.codinghorror.com/blog/archives/001018.html"""
+    return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", string_.lower())]
+
+
+def add_bool_arg(parser, name, default=False, help=""):
+    dest_name = name.replace("-", "_")
+    group = parser.add_mutually_exclusive_group(required=False)
+    group.add_argument("--" + name, dest=dest_name, action="store_true", help=help)
+    group.add_argument("--no-" + name, dest=dest_name, action="store_false", help=help)
+    parser.set_defaults(**{dest_name: default})
+
+
+def cumulative_score(pred_ages, gt_ages, L, tol=1e-6):
+    n = pred_ages.shape[0]
+    num_correct = torch.sum(torch.abs(pred_ages - gt_ages) <= L + tol)
+    cs_score = num_correct / n
+    return cs_score
+
+
+def cumulative_error(pred_ages, gt_ages, L, tol=1e-6):
+    n = pred_ages.shape[0]
+    num_correct = torch.sum(torch.abs(pred_ages - gt_ages) >= L + tol)
+    cs_score = num_correct / n
+    return cs_score
+
+
+class ParseKwargs(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        kw = {}
+        for value in values:
+            key, value = value.split("=")
+            try:
+                kw[key] = ast.literal_eval(value)
+            except ValueError:
+                kw[key] = str(value)  # fallback to string (avoid need to escape on command line)
+        setattr(namespace, self.dest, kw)
+
+
+def box_iou(box1, box2, over_second=False):
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+    If over_second == True, return mean(intersection-over-union, (inter / area2))
+
+    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+
+    Arguments:
+        box1 (Tensor[N, 4])
+        box2 (Tensor[M, 4])
+    Returns:
+        iou (Tensor[N, M]): the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+    """
+
+    def box_area(box):
+        # box = 4xn
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area1 = box_area(box1.T)
+    area2 = box_area(box2.T)
+
+    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
+    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
+
+    iou = inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)
+    if over_second:
+        return (inter / area2 + iou) / 2  # mean(inter / area2, iou)
+    else:
+        return iou
+
+
+def split_batch(bs: int, dev: int) -> Tuple[int, int]:
+    full_bs = (bs // dev) * dev
+    part_bs = bs - full_bs
+    return full_bs, part_bs
+
+
+def assign_faces(
+    persons_bboxes: List[torch.tensor], faces_bboxes: List[torch.tensor], iou_thresh: float = 0.0001
+) -> Tuple[List[Optional[int]], List[int]]:
+    """
+    Assign person to each face if it is possible.
+    Return:
+        - assigned_faces List[Optional[int]]: mapping of face_ind to person_ind
+                                            ( assigned_faces[face_ind] = person_ind ). person_ind can be None
+        - unassigned_persons_inds List[int]: persons indexes without any assigned face
+    """
+
+    assigned_faces: List[Optional[int]] = [None for _ in range(len(faces_bboxes))]
+    unassigned_persons_inds: List[int] = [p_ind for p_ind in range(len(persons_bboxes))]
+
+    if len(persons_bboxes) == 0 or len(faces_bboxes) == 0:
+        return assigned_faces, unassigned_persons_inds
+
+    cost_matrix = box_iou(torch.stack(persons_bboxes), torch.stack(faces_bboxes), over_second=True).cpu().numpy()
+    persons_indexes, face_indexes = [], []
+
+    if len(cost_matrix) > 0:
+        persons_indexes, face_indexes = linear_sum_assignment(cost_matrix, maximize=True)
+
+    matched_persons = set()
+    for person_idx, face_idx in zip(persons_indexes, face_indexes):
+        ciou = cost_matrix[person_idx][face_idx]
+        if ciou > iou_thresh:
+            if person_idx in matched_persons:
+                # Person can not be assigned twice, in reality this should not happen
+                continue
+            assigned_faces[face_idx] = person_idx
+            matched_persons.add(person_idx)
+
+    unassigned_persons_inds = [p_ind for p_ind in range(len(persons_bboxes)) if p_ind not in matched_persons]
+
+    return assigned_faces, unassigned_persons_inds
+
+
+def class_letterbox(im, new_shape=(640, 640), color=(0, 0, 0), scaleup=True):
+    # Resize and pad image while meeting stride-multiple constraints
+    shape = im.shape[:2]  # current shape [height, width]
+    if isinstance(new_shape, int):
+        new_shape = (new_shape, new_shape)
+
+    if im.shape[0] == new_shape[0] and im.shape[1] == new_shape[1]:
+        return im
+
+    # Scale ratio (new / old)
+    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+    if not scaleup:  # only scale down, do not scale up (for better val mAP)
+        r = min(r, 1.0)
+
+    # Compute padding
+    # ratio = r, r  # width, height ratios
+    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+
+    dw /= 2  # divide padding into 2 sides
+    dh /= 2
+
+    if shape[::-1] != new_unpad:  # resize
+        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
+    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
+    return im
+
+
+def prepare_classification_images(
+    img_list: List[Optional[np.ndarray]],
+    target_size: int = 224,
+    mean=IMAGENET_DEFAULT_MEAN,
+    std=IMAGENET_DEFAULT_STD,
+    device=None,
+) -> torch.tensor:
+
+    prepared_images: List[torch.tensor] = []
+
+    for img in img_list:
+        if img is None:
+            img = torch.zeros((3, target_size, target_size), dtype=torch.float32)
+            img = F.normalize(img, mean=mean, std=std)
+            img = img.unsqueeze(0)
+            prepared_images.append(img)
+            continue
+        img = class_letterbox(img, new_shape=(target_size, target_size))
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        img = img / 255.0
+        img = (img - mean) / std
+        img = img.astype(dtype=np.float32)
+
+        img = img.transpose((2, 0, 1))
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img)
+        img = img.unsqueeze(0)
+
+        prepared_images.append(img)
+
+    prepared_input = torch.concat(prepared_images)
+
+    if device:
+        prepared_input = prepared_input.to(device)
+
+    return prepared_input
+
+
+def IOU(bb1: Union[tuple, list], bb2: Union[tuple, list], norm_second_bbox: bool = False) -> float:
+    # expects [ymin, xmin, ymax, xmax], doesnt matter absolute or relative
+    assert bb1[1] < bb1[3]
+    assert bb1[0] < bb1[2]
+    assert bb2[1] < bb2[3]
+    assert bb2[0] < bb2[2]
+
+    # determine the coordinates of the intersection rectangle
+    x_left = max(bb1[1], bb2[1])
+    y_top = max(bb1[0], bb2[0])
+    x_right = min(bb1[3], bb2[3])
+    y_bottom = min(bb1[2], bb2[2])
+
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
+
+    # The intersection of two axis-aligned bounding boxes is always an
+    # axis-aligned bounding box
+    intersection_area = (x_right - x_left) * (y_bottom - y_top)
+    # compute the area of both AABBs
+    bb1_area = (bb1[3] - bb1[1]) * (bb1[2] - bb1[0])
+    bb2_area = (bb2[3] - bb2[1]) * (bb2[2] - bb2[0])
+    if not norm_second_bbox:
+        # compute the intersection over union by taking the intersection
+        # area and dividing it by the sum of prediction + ground-truth
+        # areas - the interesection area
+        iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
+    else:
+        # for cases when we search if second bbox is inside first one
+        iou = intersection_area / float(bb2_area)
+
+    assert iou >= 0.0
+    assert iou <= 1.01
+
+    return iou

mivolo/model/create_timm_model.py

@@ -0,0 +1,107 @@
+"""
+Code adapted from timm https://github.com/huggingface/pytorch-image-models
+
+Modifications and additions for mivolo by / Copyright 2023, Irina Tolstykh, Maxim Kuprashevich
+"""
+
+import os
+from typing import Any, Dict, Optional, Union
+
+import timm
+
+# register new models
+from mivolo.model.mivolo_model import *  # noqa: F403, F401
+from timm.layers import set_layer_config
+from timm.models._factory import parse_model_name
+from timm.models._helpers import load_state_dict, remap_checkpoint
+from timm.models._hub import load_model_config_from_hf
+from timm.models._pretrained import PretrainedCfg, split_model_name_tag
+from timm.models._registry import is_model, model_entrypoint
+
+
+def load_checkpoint(
+    model, checkpoint_path, use_ema=True, strict=True, remap=False, filter_keys=None, state_dict_map=None
+):
+    if os.path.splitext(checkpoint_path)[-1].lower() in (".npz", ".npy"):
+        # numpy checkpoint, try to load via model specific load_pretrained fn
+        if hasattr(model, "load_pretrained"):
+            timm.models._model_builder.load_pretrained(checkpoint_path)
+        else:
+            raise NotImplementedError("Model cannot load numpy checkpoint")
+        return
+    state_dict = load_state_dict(checkpoint_path, use_ema)
+    if remap:
+        state_dict = remap_checkpoint(model, state_dict)
+    if filter_keys:
+        for sd_key in list(state_dict.keys()):
+            for filter_key in filter_keys:
+                if filter_key in sd_key:
+                    if sd_key in state_dict:
+                        del state_dict[sd_key]
+
+    rep = []
+    if state_dict_map is not None:
+        # 'patch_embed.conv1.' : 'patch_embed.conv.'
+        for state_k in list(state_dict.keys()):
+            for target_k, target_v in state_dict_map.items():
+                if target_v in state_k:
+                    target_name = state_k.replace(target_v, target_k)
+                    state_dict[target_name] = state_dict[state_k]
+                    rep.append(state_k)
+        for r in rep:
+            if r in state_dict:
+                del state_dict[r]
+
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict if filter_keys is None else False)
+    return incompatible_keys
+
+
+def create_model(
+    model_name: str,
+    pretrained: bool = False,
+    pretrained_cfg: Optional[Union[str, Dict[str, Any], PretrainedCfg]] = None,
+    pretrained_cfg_overlay: Optional[Dict[str, Any]] = None,
+    checkpoint_path: str = "",
+    scriptable: Optional[bool] = None,
+    exportable: Optional[bool] = None,
+    no_jit: Optional[bool] = None,
+    filter_keys=None,
+    state_dict_map=None,
+    **kwargs,
+):
+    """Create a model
+    Lookup model's entrypoint function and pass relevant args to create a new model.
+    """
+    # Parameters that aren't supported by all models or are intended to only override model defaults if set
+    # should default to None in command line args/cfg. Remove them if they are present and not set so that
+    # non-supporting models don't break and default args remain in effect.
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+    model_source, model_name = parse_model_name(model_name)
+    if model_source == "hf-hub":
+        assert not pretrained_cfg, "pretrained_cfg should not be set when sourcing model from Hugging Face Hub."
+        # For model names specified in the form `hf-hub:path/architecture_name@revision`,
+        # load model weights + pretrained_cfg from Hugging Face hub.
+        pretrained_cfg, model_name = load_model_config_from_hf(model_name)
+    else:
+        model_name, pretrained_tag = split_model_name_tag(model_name)
+        if not pretrained_cfg:
+            # a valid pretrained_cfg argument takes priority over tag in model name
+            pretrained_cfg = pretrained_tag
+
+    if not is_model(model_name):
+        raise RuntimeError("Unknown model (%s)" % model_name)
+
+    create_fn = model_entrypoint(model_name)
+    with set_layer_config(scriptable=scriptable, exportable=exportable, no_jit=no_jit):
+        model = create_fn(
+            pretrained=pretrained,
+            pretrained_cfg=pretrained_cfg,
+            pretrained_cfg_overlay=pretrained_cfg_overlay,
+            **kwargs,
+        )
+
+    if checkpoint_path:
+        load_checkpoint(model, checkpoint_path, filter_keys=filter_keys, state_dict_map=state_dict_map)
+
+    return model

mivolo/model/cross_bottleneck_attn.py

@@ -0,0 +1,116 @@
+"""
+Code based on timm https://github.com/huggingface/pytorch-image-models
+
+Modifications and additions for mivolo by / Copyright 2023, Irina Tolstykh, Maxim Kuprashevich
+"""
+
+import torch
+import torch.nn as nn
+from timm.layers.bottleneck_attn import PosEmbedRel
+from timm.layers.helpers import make_divisible
+from timm.layers.mlp import Mlp
+from timm.layers.trace_utils import _assert
+from timm.layers.weight_init import trunc_normal_
+
+
+class CrossBottleneckAttn(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out=None,
+        feat_size=None,
+        stride=1,
+        num_heads=4,
+        dim_head=None,
+        qk_ratio=1.0,
+        qkv_bias=False,
+        scale_pos_embed=False,
+    ):
+        super().__init__()
+        assert feat_size is not None, "A concrete feature size matching expected input (H, W) is required"
+        dim_out = dim_out or dim
+        assert dim_out % num_heads == 0
+
+        self.num_heads = num_heads
+        self.dim_head_qk = dim_head or make_divisible(dim_out * qk_ratio, divisor=8) // num_heads
+        self.dim_head_v = dim_out // self.num_heads
+        self.dim_out_qk = num_heads * self.dim_head_qk
+        self.dim_out_v = num_heads * self.dim_head_v
+        self.scale = self.dim_head_qk**-0.5
+        self.scale_pos_embed = scale_pos_embed
+
+        self.qkv_f = nn.Conv2d(dim, self.dim_out_qk * 2 + self.dim_out_v, 1, bias=qkv_bias)
+        self.qkv_p = nn.Conv2d(dim, self.dim_out_qk * 2 + self.dim_out_v, 1, bias=qkv_bias)
+
+        # NOTE I'm only supporting relative pos embedding for now
+        self.pos_embed = PosEmbedRel(feat_size, dim_head=self.dim_head_qk, scale=self.scale)
+
+        self.norm = nn.LayerNorm([self.dim_out_v * 2, *feat_size])
+        mlp_ratio = 4
+        self.mlp = Mlp(
+            in_features=self.dim_out_v * 2,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=nn.GELU,
+            out_features=dim_out,
+            drop=0,
+            use_conv=True,
+        )
+
+        self.pool = nn.AvgPool2d(2, 2) if stride == 2 else nn.Identity()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        trunc_normal_(self.qkv_f.weight, std=self.qkv_f.weight.shape[1] ** -0.5)  # fan-in
+        trunc_normal_(self.qkv_p.weight, std=self.qkv_p.weight.shape[1] ** -0.5)  # fan-in
+        trunc_normal_(self.pos_embed.height_rel, std=self.scale)
+        trunc_normal_(self.pos_embed.width_rel, std=self.scale)
+
+    def get_qkv(self, x, qvk_conv):
+        B, C, H, W = x.shape
+
+        x = qvk_conv(x)  # B, (2 * dim_head_qk + dim_head_v) * num_heads, H, W
+
+        q, k, v = torch.split(x, [self.dim_out_qk, self.dim_out_qk, self.dim_out_v], dim=1)
+
+        q = q.reshape(B * self.num_heads, self.dim_head_qk, -1).transpose(-1, -2)
+        k = k.reshape(B * self.num_heads, self.dim_head_qk, -1)  # no transpose, for q @ k
+        v = v.reshape(B * self.num_heads, self.dim_head_v, -1).transpose(-1, -2)
+
+        return q, k, v
+
+    def apply_attn(self, q, k, v, B, H, W, dropout=None):
+        if self.scale_pos_embed:
+            attn = (q @ k + self.pos_embed(q)) * self.scale  # B * num_heads, H * W, H * W
+        else:
+            attn = (q @ k) * self.scale + self.pos_embed(q)
+        attn = attn.softmax(dim=-1)
+        if dropout:
+            attn = dropout(attn)
+
+        out = (attn @ v).transpose(-1, -2).reshape(B, self.dim_out_v, H, W)  # B, dim_out, H, W
+        return out
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+
+        dim = int(C / 2)
+        x1 = x[:, :dim, :, :]
+        x2 = x[:, dim:, :, :]
+
+        _assert(H == self.pos_embed.height, "")
+        _assert(W == self.pos_embed.width, "")
+
+        q_f, k_f, v_f = self.get_qkv(x1, self.qkv_f)
+        q_p, k_p, v_p = self.get_qkv(x2, self.qkv_p)
+
+        # person to face
+        out_f = self.apply_attn(q_f, k_p, v_p, B, H, W)
+        # face to person
+        out_p = self.apply_attn(q_p, k_f, v_f, B, H, W)
+
+        x_pf = torch.cat((out_f, out_p), dim=1)  # B, dim_out * 2, H, W
+        x_pf = self.norm(x_pf)
+        x_pf = self.mlp(x_pf)  # B, dim_out, H, W
+
+        out = self.pool(x_pf)
+        return out

mivolo/model/mi_volo.py

@@ -0,0 +1,229 @@
+import logging
+from typing import Optional
+
+import numpy as np
+import torch
+from mivolo.data.misc import prepare_classification_images
+from mivolo.model.create_timm_model import create_model
+from mivolo.structures import PersonAndFaceCrops, PersonAndFaceResult
+from timm.data import resolve_data_config
+
+_logger = logging.getLogger("MiVOLO")
+has_compile = hasattr(torch, "compile")
+
+
+class Meta:
+    def __init__(self):
+        self.min_age = None
+        self.max_age = None
+        self.avg_age = None
+        self.num_classes = None
+
+        self.in_chans = 3
+        self.with_persons_model = False
+        self.disable_faces = False
+        self.use_persons = True
+        self.only_age = False
+
+        self.num_classes_gender = 2
+
+    def load_from_ckpt(self, ckpt_path: str, disable_faces: bool = False, use_persons: bool = True) -> "Meta":
+
+        state = torch.load(ckpt_path, map_location="cpu")
+
+        self.min_age = state["min_age"]
+        self.max_age = state["max_age"]
+        self.avg_age = state["avg_age"]
+        self.only_age = state["no_gender"]
+
+        only_age = state["no_gender"]
+
+        self.disable_faces = disable_faces
+        if "with_persons_model" in state:
+            self.with_persons_model = state["with_persons_model"]
+        else:
+            self.with_persons_model = True if "patch_embed.conv1.0.weight" in state["state_dict"] else False
+
+        self.num_classes = 1 if only_age else 3
+        self.in_chans = 3 if not self.with_persons_model else 6
+        self.use_persons = use_persons and self.with_persons_model
+
+        if not self.with_persons_model and self.disable_faces:
+            raise ValueError("You can not use disable-faces for faces-only model")
+        if self.with_persons_model and self.disable_faces and not self.use_persons:
+            raise ValueError("You can not disable faces and persons together")
+
+        return self
+
+    def __str__(self):
+        attrs = vars(self)
+        attrs.update({"use_person_crops": self.use_person_crops, "use_face_crops": self.use_face_crops})
+        return ", ".join("%s: %s" % item for item in attrs.items())
+
+    @property
+    def use_person_crops(self) -> bool:
+        return self.with_persons_model and self.use_persons
+
+    @property
+    def use_face_crops(self) -> bool:
+        return not self.disable_faces or not self.with_persons_model
+
+
+class MiVOLO:
+    def __init__(
+        self,
+        ckpt_path: str,
+        device: str = "cpu",
+        half: bool = True,
+        disable_faces: bool = False,
+        use_persons: bool = True,
+        verbose: bool = False,
+        torchcompile: Optional[str] = None,
+    ):
+        self.verbose = verbose
+        self.device = torch.device(device)
+        self.half = half and self.device.type != "cpu"
+
+        self.meta: Meta = Meta().load_from_ckpt(ckpt_path, disable_faces, use_persons)
+        if self.verbose:
+            _logger.info(f"Model meta:\n{str(self.meta)}")
+
+        model_name = "mivolo_d1_224"
+        self.model = create_model(
+            model_name=model_name,
+            num_classes=self.meta.num_classes,
+            in_chans=self.meta.in_chans,
+            pretrained=False,
+            checkpoint_path=ckpt_path,
+            filter_keys=["fds."],
+        )
+        self.param_count = sum([m.numel() for m in self.model.parameters()])
+        _logger.info(f"Model {model_name} created, param count: {self.param_count}")
+
+        self.data_config = resolve_data_config(
+            model=self.model,
+            verbose=verbose,
+            use_test_size=True,
+        )
+        self.data_config["crop_pct"] = 1.0
+        c, h, w = self.data_config["input_size"]
+        assert h == w, "Incorrect data_config"
+        self.input_size = w
+
+        self.model = self.model.to(self.device)
+
+        if torchcompile:
+            assert has_compile, "A version of torch w/ torch.compile() is required for --compile, possibly a nightly."
+            torch._dynamo.reset()
+            self.model = torch.compile(self.model, backend=torchcompile)
+
+        self.model.eval()
+        if self.half:
+            self.model = self.model.half()
+
+    def warmup(self, batch_size: int, steps=10):
+        if self.meta.with_persons_model:
+            input_size = (6, self.input_size, self.input_size)
+        else:
+            input_size = self.data_config["input_size"]
+
+        input = torch.randn((batch_size,) + tuple(input_size)).to(self.device)
+
+        for _ in range(steps):
+            out = self.inference(input)  # noqa: F841
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def inference(self, model_input: torch.tensor) -> torch.tensor:
+
+        with torch.no_grad():
+            if self.half:
+                model_input = model_input.half()
+            output = self.model(model_input)
+        return output
+
+    def predict(self, image: np.ndarray, detected_bboxes: PersonAndFaceResult):
+        if detected_bboxes.n_objects == 0:
+            return
+
+        faces_input, person_input, faces_inds, bodies_inds = self.prepare_crops(image, detected_bboxes)
+
+        if self.meta.with_persons_model:
+            model_input = torch.cat((faces_input, person_input), dim=1)
+        else:
+            model_input = faces_input
+        output = self.inference(model_input)
+
+        # write gender and age results into detected_bboxes
+        self.fill_in_results(output, detected_bboxes, faces_inds, bodies_inds)
+
+    def fill_in_results(self, output, detected_bboxes, faces_inds, bodies_inds):
+        if self.meta.only_age:
+            age_output = output
+            gender_probs, gender_indx = None, None
+        else:
+            age_output = output[:, 2]
+            gender_output = output[:, :2].softmax(-1)
+            gender_probs, gender_indx = gender_output.topk(1)
+
+        assert output.shape[0] == len(faces_inds) == len(bodies_inds)
+
+        # per face
+        for index in range(output.shape[0]):
+            face_ind = faces_inds[index]
+            body_ind = bodies_inds[index]
+
+            # get_age
+            age = age_output[index].item()
+            age = age * (self.meta.max_age - self.meta.min_age) + self.meta.avg_age
+            age = round(age, 2)
+
+            detected_bboxes.set_age(face_ind, age)
+            detected_bboxes.set_age(body_ind, age)
+
+            _logger.info(f"\tage: {age}")
+
+            if gender_probs is not None:
+                gender = "male" if gender_indx[index].item() == 0 else "female"
+                gender_score = gender_probs[index].item()
+
+                _logger.info(f"\tgender: {gender} [{int(gender_score * 100)}%]")
+
+                detected_bboxes.set_gender(face_ind, gender, gender_score)
+                detected_bboxes.set_gender(body_ind, gender, gender_score)
+
+    def prepare_crops(self, image: np.ndarray, detected_bboxes: PersonAndFaceResult):
+
+        if self.meta.use_person_crops and self.meta.use_face_crops:
+            detected_bboxes.associate_faces_with_persons()
+
+        crops: PersonAndFaceCrops = detected_bboxes.collect_crops(image)
+        (bodies_inds, bodies_crops), (faces_inds, faces_crops) = crops.get_faces_with_bodies(
+            self.meta.use_person_crops, self.meta.use_face_crops
+        )
+
+        if not self.meta.use_face_crops:
+            assert all(f is None for f in faces_crops)
+
+        faces_input = prepare_classification_images(
+            faces_crops, self.input_size, self.data_config["mean"], self.data_config["std"], device=self.device
+        )
+
+        if not self.meta.use_person_crops:
+            assert all(p is None for p in bodies_crops)
+
+        person_input = prepare_classification_images(
+            bodies_crops, self.input_size, self.data_config["mean"], self.data_config["std"], device=self.device
+        )
+
+        _logger.info(
+            f"faces_input: {faces_input.shape if faces_input is not None else None}, "
+            f"person_input: {person_input.shape if person_input is not None else None}"
+        )
+
+        return faces_input, person_input, faces_inds, bodies_inds
+
+
+if __name__ == "__main__":
+    model = MiVOLO("../pretrained/checkpoint-377.pth.tar", half=True, device="cuda:0")

mivolo/model/mivolo_model.py

@@ -0,0 +1,402 @@
+"""
+Code adapted from timm https://github.com/huggingface/pytorch-image-models
+
+Modifications and additions for mivolo by / Copyright 2023, Irina Tolstykh, Maxim Kuprashevich
+"""
+
+import torch
+import torch.nn as nn
+from mivolo.model.cross_bottleneck_attn import CrossBottleneckAttn
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.layers import trunc_normal_
+from timm.models._builder import build_model_with_cfg
+from timm.models._registry import register_model
+from timm.models.volo import VOLO
+
+__all__ = ["MiVOLOModel"]  # model_registry will add each entrypoint fn to this
+
+
+def _cfg(url="", **kwargs):
+    return {
+        "url": url,
+        "num_classes": 1000,
+        "input_size": (3, 224, 224),
+        "pool_size": None,
+        "crop_pct": 0.96,
+        "interpolation": "bicubic",
+        "fixed_input_size": True,
+        "mean": IMAGENET_DEFAULT_MEAN,
+        "std": IMAGENET_DEFAULT_STD,
+        "first_conv": None,
+        "classifier": ("head", "aux_head"),
+        **kwargs,
+    }
+
+
+default_cfgs = {
+    "mivolo_d1_224": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d1_224_84.2.pth.tar", crop_pct=0.96
+    ),
+    "mivolo_d1_384": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d1_384_85.2.pth.tar",
+        crop_pct=1.0,
+        input_size=(3, 384, 384),
+    ),
+    "mivolo_d2_224": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d2_224_85.2.pth.tar", crop_pct=0.96
+    ),
+    "mivolo_d2_384": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d2_384_86.0.pth.tar",
+        crop_pct=1.0,
+        input_size=(3, 384, 384),
+    ),
+    "mivolo_d3_224": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d3_224_85.4.pth.tar", crop_pct=0.96
+    ),
+    "mivolo_d3_448": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d3_448_86.3.pth.tar",
+        crop_pct=1.0,
+        input_size=(3, 448, 448),
+    ),
+    "mivolo_d4_224": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d4_224_85.7.pth.tar", crop_pct=0.96
+    ),
+    "mivolo_d4_448": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d4_448_86.79.pth.tar",
+        crop_pct=1.15,
+        input_size=(3, 448, 448),
+    ),
+    "mivolo_d5_224": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d5_224_86.10.pth.tar", crop_pct=0.96
+    ),
+    "mivolo_d5_448": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d5_448_87.0.pth.tar",
+        crop_pct=1.15,
+        input_size=(3, 448, 448),
+    ),
+    "mivolo_d5_512": _cfg(
+        url="https://github.com/sail-sg/volo/releases/download/volo_1/d5_512_87.07.pth.tar",
+        crop_pct=1.15,
+        input_size=(3, 512, 512),
+    ),
+}
+
+
+def get_output_size(input_shape, conv_layer):
+    padding = conv_layer.padding
+    dilation = conv_layer.dilation
+    kernel_size = conv_layer.kernel_size
+    stride = conv_layer.stride
+
+    output_size = [
+        ((input_shape[i] + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) // stride[i]) + 1 for i in range(2)
+    ]
+    return output_size
+
+
+def get_output_size_module(input_size, stem):
+    output_size = input_size
+
+    for module in stem:
+        if isinstance(module, nn.Conv2d):
+            output_size = [
+                (
+                    (output_size[i] + 2 * module.padding[i] - module.dilation[i] * (module.kernel_size[i] - 1) - 1)
+                    // module.stride[i]
+                )
+                + 1
+                for i in range(2)
+            ]
+
+    return output_size
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding."""
+
+    def __init__(
+        self, img_size=224, stem_conv=False, stem_stride=1, patch_size=8, in_chans=3, hidden_dim=64, embed_dim=384
+    ):
+        super().__init__()
+        assert patch_size in [4, 8, 16]
+        assert in_chans in [3, 6]
+        self.with_persons_model = in_chans == 6
+        self.use_cross_attn = True
+
+        if stem_conv:
+            if not self.with_persons_model:
+                self.conv = self.create_stem(stem_stride, in_chans, hidden_dim)
+            else:
+                self.conv = True  # just to match interface
+                # split
+                self.conv1 = self.create_stem(stem_stride, 3, hidden_dim)
+                self.conv2 = self.create_stem(stem_stride, 3, hidden_dim)
+        else:
+            self.conv = None
+
+        if self.with_persons_model:
+
+            self.proj1 = nn.Conv2d(
+                hidden_dim, embed_dim, kernel_size=patch_size // stem_stride, stride=patch_size // stem_stride
+            )
+            self.proj2 = nn.Conv2d(
+                hidden_dim, embed_dim, kernel_size=patch_size // stem_stride, stride=patch_size // stem_stride
+            )
+
+            stem_out_shape = get_output_size_module((img_size, img_size), self.conv1)
+            self.proj_output_size = get_output_size(stem_out_shape, self.proj1)
+
+            self.map = CrossBottleneckAttn(embed_dim, dim_out=embed_dim, num_heads=1, feat_size=self.proj_output_size)
+
+        else:
+            self.proj = nn.Conv2d(
+                hidden_dim, embed_dim, kernel_size=patch_size // stem_stride, stride=patch_size // stem_stride
+            )
+
+        self.patch_dim = img_size // patch_size
+        self.num_patches = self.patch_dim**2
+
+    def create_stem(self, stem_stride, in_chans, hidden_dim):
+        return nn.Sequential(
+            nn.Conv2d(in_chans, hidden_dim, kernel_size=7, stride=stem_stride, padding=3, bias=False),  # 112x112
+            nn.BatchNorm2d(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1, bias=False),  # 112x112
+            nn.BatchNorm2d(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1, bias=False),  # 112x112
+            nn.BatchNorm2d(hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        if self.conv is not None:
+            if self.with_persons_model:
+                x1 = x[:, :3]
+                x2 = x[:, 3:]
+
+                x1 = self.conv1(x1)
+                x1 = self.proj1(x1)
+
+                x2 = self.conv2(x2)
+                x2 = self.proj2(x2)
+
+                x = torch.cat([x1, x2], dim=1)
+                x = self.map(x)
+            else:
+                x = self.conv(x)
+                x = self.proj(x)  # B, C, H, W
+
+        return x
+
+
+class MiVOLOModel(VOLO):
+    """
+    Vision Outlooker, the main class of our model
+    """
+
+    def __init__(
+        self,
+        layers,
+        img_size=224,
+        in_chans=3,
+        num_classes=1000,
+        global_pool="token",
+        patch_size=8,
+        stem_hidden_dim=64,
+        embed_dims=None,
+        num_heads=None,
+        downsamples=(True, False, False, False),
+        outlook_attention=(True, False, False, False),
+        mlp_ratio=3.0,
+        qkv_bias=False,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        post_layers=("ca", "ca"),
+        use_aux_head=True,
+        use_mix_token=False,
+        pooling_scale=2,
+    ):
+        super().__init__(
+            layers,
+            img_size,
+            in_chans,
+            num_classes,
+            global_pool,
+            patch_size,
+            stem_hidden_dim,
+            embed_dims,
+            num_heads,
+            downsamples,
+            outlook_attention,
+            mlp_ratio,
+            qkv_bias,
+            drop_rate,
+            attn_drop_rate,
+            drop_path_rate,
+            norm_layer,
+            post_layers,
+            use_aux_head,
+            use_mix_token,
+            pooling_scale,
+        )
+
+        self.patch_embed = PatchEmbed(
+            stem_conv=True,
+            stem_stride=2,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            hidden_dim=stem_hidden_dim,
+            embed_dim=embed_dims[0],
+        )
+
+        trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(self._init_weights)
+
+    def forward_features(self, x):
+        x = self.patch_embed(x).permute(0, 2, 3, 1)  # B,C,H,W-> B,H,W,C
+
+        # step2: tokens learning in the two stages
+        x = self.forward_tokens(x)
+
+        # step3: post network, apply class attention or not
+        if self.post_network is not None:
+            x = self.forward_cls(x)
+        x = self.norm(x)
+        return x
+
+    def forward_head(self, x, pre_logits: bool = False, targets=None, epoch=None):
+        if self.global_pool == "avg":
+            out = x.mean(dim=1)
+        elif self.global_pool == "token":
+            out = x[:, 0]
+        else:
+            out = x
+        if pre_logits:
+            return out
+
+        features = out
+        fds_enabled = hasattr(self, "_fds_forward")
+        if fds_enabled:
+            features = self._fds_forward(features, targets, epoch)
+
+        out = self.head(features)
+        if self.aux_head is not None:
+            # generate classes in all feature tokens, see token labeling
+            aux = self.aux_head(x[:, 1:])
+            out = out + 0.5 * aux.max(1)[0]
+
+        return (out, features) if (fds_enabled and self.training) else out
+
+    def forward(self, x, targets=None, epoch=None):
+        """simplified forward (without mix token training)"""
+        x = self.forward_features(x)
+        x = self.forward_head(x, targets=targets, epoch=epoch)
+        return x
+
+
+def _create_mivolo(variant, pretrained=False, **kwargs):
+    if kwargs.get("features_only", None):
+        raise RuntimeError("features_only not implemented for Vision Transformer models.")
+    return build_model_with_cfg(MiVOLOModel, variant, pretrained, **kwargs)
+
+
+@register_model
+def mivolo_d1_224(pretrained=False, **kwargs):
+    model_args = dict(layers=(4, 4, 8, 2), embed_dims=(192, 384, 384, 384), num_heads=(6, 12, 12, 12), **kwargs)
+    model = _create_mivolo("mivolo_d1_224", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d1_384(pretrained=False, **kwargs):
+    model_args = dict(layers=(4, 4, 8, 2), embed_dims=(192, 384, 384, 384), num_heads=(6, 12, 12, 12), **kwargs)
+    model = _create_mivolo("mivolo_d1_384", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d2_224(pretrained=False, **kwargs):
+    model_args = dict(layers=(6, 4, 10, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d2_224", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d2_384(pretrained=False, **kwargs):
+    model_args = dict(layers=(6, 4, 10, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d2_384", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d3_224(pretrained=False, **kwargs):
+    model_args = dict(layers=(8, 8, 16, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d3_224", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d3_448(pretrained=False, **kwargs):
+    model_args = dict(layers=(8, 8, 16, 4), embed_dims=(256, 512, 512, 512), num_heads=(8, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d3_448", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d4_224(pretrained=False, **kwargs):
+    model_args = dict(layers=(8, 8, 16, 4), embed_dims=(384, 768, 768, 768), num_heads=(12, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d4_224", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d4_448(pretrained=False, **kwargs):
+    """VOLO-D4 model, Params: 193M"""
+    model_args = dict(layers=(8, 8, 16, 4), embed_dims=(384, 768, 768, 768), num_heads=(12, 16, 16, 16), **kwargs)
+    model = _create_mivolo("mivolo_d4_448", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d5_224(pretrained=False, **kwargs):
+    model_args = dict(
+        layers=(12, 12, 20, 4),
+        embed_dims=(384, 768, 768, 768),
+        num_heads=(12, 16, 16, 16),
+        mlp_ratio=4,
+        stem_hidden_dim=128,
+        **kwargs
+    )
+    model = _create_mivolo("mivolo_d5_224", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d5_448(pretrained=False, **kwargs):
+    model_args = dict(
+        layers=(12, 12, 20, 4),
+        embed_dims=(384, 768, 768, 768),
+        num_heads=(12, 16, 16, 16),
+        mlp_ratio=4,
+        stem_hidden_dim=128,
+        **kwargs
+    )
+    model = _create_mivolo("mivolo_d5_448", pretrained=pretrained, **model_args)
+    return model
+
+
+@register_model
+def mivolo_d5_512(pretrained=False, **kwargs):
+    model_args = dict(
+        layers=(12, 12, 20, 4),
+        embed_dims=(384, 768, 768, 768),
+        num_heads=(12, 16, 16, 16),
+        mlp_ratio=4,
+        stem_hidden_dim=128,
+        **kwargs
+    )
+    model = _create_mivolo("mivolo_d5_512", pretrained=pretrained, **model_args)
+    return model

mivolo/model/yolo_detector.py

@@ -0,0 +1,48 @@
+import os
+from typing import Dict, Union
+
+import numpy as np
+import PIL
+import torch
+from mivolo.structures import PersonAndFaceResult
+from ultralytics import YOLO
+# from ultralytics.yolo.engine.results import Results
+
+# because of ultralytics bug it is important to unset CUBLAS_WORKSPACE_CONFIG after the module importing
+os.unsetenv("CUBLAS_WORKSPACE_CONFIG")
+
+
+class Detector:
+    def __init__(
+        self,
+        weights: str,
+        device: str = "cpu",
+        half: bool = True,
+        verbose: bool = False,
+        conf_thresh: float = 0.4,
+        iou_thresh: float = 0.7,
+    ):
+        self.yolo = YOLO(weights)
+        self.yolo.fuse()
+
+        self.device = torch.device(device)
+        self.half = half and self.device.type != "cpu"
+
+        if self.half:
+            self.yolo.model = self.yolo.model.half()
+
+        self.detector_names: Dict[int, str] = self.yolo.model.names
+
+        # init yolo.predictor
+        self.detector_kwargs = {
+            "conf": conf_thresh, "iou": iou_thresh, "half": self.half, "verbose": verbose}
+        # self.yolo.predict(**self.detector_kwargs)
+
+    def predict(self, image: Union[np.ndarray, str, "PIL.Image"]) -> PersonAndFaceResult:
+        results = self.yolo.predict(image, **self.detector_kwargs)[0]
+        return PersonAndFaceResult(results)
+
+    def track(self, image: Union[np.ndarray, str, "PIL.Image"]) -> PersonAndFaceResult:
+        results = self.yolo.track(
+            image, persist=True, **self.detector_kwargs)[0]
+        return PersonAndFaceResult(results)

mivolo/predictor.py

@@ -0,0 +1,68 @@
+from collections import defaultdict
+from typing import Dict, Generator, List, Optional, Tuple
+
+import cv2
+import numpy as np
+import tqdm
+from mivolo.model.mi_volo import MiVOLO
+from mivolo.model.yolo_detector import Detector
+from mivolo.structures import AGE_GENDER_TYPE, PersonAndFaceResult
+
+
+class Predictor:
+    def __init__(self, config, verbose: bool = False):
+        self.detector = Detector(config.detector_weights, config.device, verbose=verbose)
+        self.age_gender_model = MiVOLO(
+            config.checkpoint,
+            config.device,
+            half=True,
+            use_persons=config.with_persons,
+            disable_faces=config.disable_faces,
+            verbose=verbose,
+        )
+        self.draw = config.draw
+
+    def recognize(self, image: np.ndarray) -> Tuple[PersonAndFaceResult, Optional[np.ndarray]]:
+        detected_objects: PersonAndFaceResult = self.detector.predict(image)
+        self.age_gender_model.predict(image, detected_objects)
+
+        out_im = None
+        if self.draw:
+            # plot results on image
+            out_im = detected_objects.plot()
+
+        return detected_objects, out_im
+
+    def recognize_video(self, source: str) -> Generator:
+        video_capture = cv2.VideoCapture(source)
+        if not video_capture.isOpened():
+            raise ValueError(f"Failed to open video source {source}")
+
+        detected_objects_history: Dict[int, List[AGE_GENDER_TYPE]] = defaultdict(list)
+
+        total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+        for _ in tqdm.tqdm(range(total_frames)):
+            ret, frame = video_capture.read()
+            if not ret:
+                break
+
+            detected_objects: PersonAndFaceResult = self.detector.track(frame)
+            self.age_gender_model.predict(frame, detected_objects)
+
+            current_frame_objs = detected_objects.get_results_for_tracking()
+            cur_persons: Dict[int, AGE_GENDER_TYPE] = current_frame_objs[0]
+            cur_faces: Dict[int, AGE_GENDER_TYPE] = current_frame_objs[1]
+
+            # add tr_persons and tr_faces to history
+            for guid, data in cur_persons.items():
+                # not useful for tracking :)
+                if None not in data: 
+                    detected_objects_history[guid].append(data)
+            for guid, data in cur_faces.items():
+                if None not in data:
+                    detected_objects_history[guid].append(data)
+
+            detected_objects.set_tracked_age_gender(detected_objects_history)
+            if self.draw:
+                frame = detected_objects.plot()
+            yield detected_objects_history, frame

mivolo/structures.py

@@ -0,0 +1,493 @@
+import math
+import os
+from copy import deepcopy
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+import torch
+from mivolo.data.misc import aggregate_votes_winsorized, assign_faces, box_iou, cropout_black_parts
+from ultralytics.engine.results import Results
+from ultralytics.utils.plotting import Annotator, colors
+
+# because of ultralytics bug it is important to unset CUBLAS_WORKSPACE_CONFIG after the module importing
+os.unsetenv("CUBLAS_WORKSPACE_CONFIG")
+
+AGE_GENDER_TYPE = Tuple[float, str]
+
+
+class PersonAndFaceCrops:
+    def __init__(self):
+        # int: index of person along results
+        self.crops_persons: Dict[int, np.ndarray] = {}
+
+        # int: index of face along results
+        self.crops_faces: Dict[int, np.ndarray] = {}
+
+        # int: index of face along results
+        self.crops_faces_wo_body: Dict[int, np.ndarray] = {}
+
+        # int: index of person along results
+        self.crops_persons_wo_face: Dict[int, np.ndarray] = {}
+
+    def _add_to_output(
+        self, crops: Dict[int, np.ndarray], out_crops: List[np.ndarray], out_crop_inds: List[Optional[int]]
+    ):
+        inds_to_add = list(crops.keys())
+        crops_to_add = list(crops.values())
+        out_crops.extend(crops_to_add)
+        out_crop_inds.extend(inds_to_add)
+
+    def _get_all_faces(
+        self, use_persons: bool, use_faces: bool
+    ) -> Tuple[List[Optional[int]], List[Optional[np.ndarray]]]:
+        """
+        Returns
+            if use_persons and use_faces
+                faces: faces_with_bodies + faces_without_bodies + [None] * len(crops_persons_wo_face)
+            if use_persons and not use_faces
+                faces: [None] * n_persons
+            if not use_persons and use_faces:
+                faces: faces_with_bodies + faces_without_bodies
+        """
+
+        def add_none_to_output(faces_inds, faces_crops, num):
+            faces_inds.extend([None for _ in range(num)])
+            faces_crops.extend([None for _ in range(num)])
+
+        faces_inds: List[Optional[int]] = []
+        faces_crops: List[Optional[np.ndarray]] = []
+
+        if not use_faces:
+            add_none_to_output(faces_inds, faces_crops, len(
+                self.crops_persons) + len(self.crops_persons_wo_face))
+            return faces_inds, faces_crops
+
+        self._add_to_output(self.crops_faces, faces_crops, faces_inds)
+        self._add_to_output(self.crops_faces_wo_body, faces_crops, faces_inds)
+
+        if use_persons:
+            add_none_to_output(faces_inds, faces_crops,
+                               len(self.crops_persons_wo_face))
+
+        return faces_inds, faces_crops
+
+    def _get_all_bodies(
+        self, use_persons: bool, use_faces: bool
+    ) -> Tuple[List[Optional[int]], List[Optional[np.ndarray]]]:
+        """
+        Returns
+            if use_persons and use_faces
+                persons: bodies_with_faces + [None] * len(faces_without_bodies) + bodies_without_faces
+            if use_persons and not use_faces
+                persons: bodies_with_faces + bodies_without_faces
+            if not use_persons and use_faces
+                persons: [None] * n_faces
+        """
+
+        def add_none_to_output(bodies_inds, bodies_crops, num):
+            bodies_inds.extend([None for _ in range(num)])
+            bodies_crops.extend([None for _ in range(num)])
+
+        bodies_inds: List[Optional[int]] = []
+        bodies_crops: List[Optional[np.ndarray]] = []
+
+        if not use_persons:
+            add_none_to_output(bodies_inds, bodies_crops, len(
+                self.crops_faces) + len(self.crops_faces_wo_body))
+            return bodies_inds, bodies_crops
+
+        self._add_to_output(self.crops_persons, bodies_crops, bodies_inds)
+        if use_faces:
+            add_none_to_output(bodies_inds, bodies_crops,
+                               len(self.crops_faces_wo_body))
+
+        self._add_to_output(self.crops_persons_wo_face,
+                            bodies_crops, bodies_inds)
+
+        return bodies_inds, bodies_crops
+
+    def get_faces_with_bodies(self, use_persons: bool, use_faces: bool):
+        """
+        Return
+            faces: faces_with_bodies, faces_without_bodies, [None] * len(crops_persons_wo_face)
+            persons: bodies_with_faces, [None] * len(faces_without_bodies), bodies_without_faces
+        """
+
+        bodies_inds, bodies_crops = self._get_all_bodies(
+            use_persons, use_faces)
+        faces_inds, faces_crops = self._get_all_faces(use_persons, use_faces)
+
+        return (bodies_inds, bodies_crops), (faces_inds, faces_crops)
+
+    def save(self, out_dir="output"):
+        ind = 0
+        os.makedirs(out_dir, exist_ok=True)
+        for crops in [self.crops_persons, self.crops_faces, self.crops_faces_wo_body, self.crops_persons_wo_face]:
+            for crop in crops.values():
+                if crop is None:
+                    continue
+                out_name = os.path.join(out_dir, f"{ind}_crop.jpg")
+                cv2.imwrite(out_name, crop)
+                ind += 1
+
+
+class PersonAndFaceResult:
+    def __init__(self, results: Results):
+
+        self.yolo_results = results
+        names = set(results.names.values())
+        assert "person" in names and "face" in names
+
+        # initially no faces and persons are associated to each other
+        self.face_to_person_map: Dict[int, Optional[int]] = {
+            ind: None for ind in self.get_bboxes_inds("face")}
+        self.unassigned_persons_inds: List[int] = self.get_bboxes_inds(
+            "person")
+        n_objects = len(self.yolo_results.boxes)
+        self.ages: List[Optional[float]] = [None for _ in range(n_objects)]
+        self.genders: List[Optional[str]] = [None for _ in range(n_objects)]
+        self.gender_scores: List[Optional[float]] = [
+            None for _ in range(n_objects)]
+
+    @property
+    def n_objects(self) -> int:
+        return len(self.yolo_results.boxes)
+
+    def get_bboxes_inds(self, category: str) -> List[int]:
+        bboxes: List[int] = []
+        for ind, det in enumerate(self.yolo_results.boxes):
+            name = self.yolo_results.names[int(det.cls)]
+            if name == category:
+                bboxes.append(ind)
+
+        return bboxes
+
+    def get_distance_to_center(self, bbox_ind: int) -> float:
+        """
+        Calculate euclidian distance between bbox center and image center.
+        """
+        im_h, im_w = self.yolo_results[bbox_ind].orig_shape
+        x1, y1, x2, y2 = self.get_bbox_by_ind(bbox_ind).cpu().numpy()
+        center_x, center_y = (x1 + x2) / 2, (y1 + y2) / 2
+        dist = math.dist([center_x, center_y], [im_w / 2, im_h / 2])
+        return dist
+
+    def plot(
+        self,
+        conf=False,
+        line_width=None,
+        font_size=None,
+        font="Arial.ttf",
+        pil=False,
+        img=None,
+        labels=True,
+        boxes=True,
+        probs=True,
+        ages=True,
+        genders=True,
+        gender_probs=False,
+    ):
+        """
+        Plots the detection results on an input RGB image. Accepts a numpy array (cv2) or a PIL Image.
+        Args:
+            conf (bool): Whether to plot the detection confidence score.
+            line_width (float, optional): The line width of the bounding boxes. If None, it is scaled to the image size.
+            font_size (float, optional): The font size of the text. If None, it is scaled to the image size.
+            font (str): The font to use for the text.
+            pil (bool): Whether to return the image as a PIL Image.
+            img (numpy.ndarray): Plot to another image. if not, plot to original image.
+            labels (bool): Whether to plot the label of bounding boxes.
+            boxes (bool): Whether to plot the bounding boxes.
+            probs (bool): Whether to plot classification probability
+            ages (bool): Whether to plot the age of bounding boxes.
+            genders (bool): Whether to plot the genders of bounding boxes.
+            gender_probs (bool): Whether to plot gender classification probability
+        Returns:
+            (numpy.ndarray): A numpy array of the annotated image.
+        """
+
+        # return self.yolo_results.plot()
+        colors_by_ind = {}
+        for face_ind, person_ind in self.face_to_person_map.items():
+            if person_ind is not None:
+                colors_by_ind[face_ind] = face_ind + 2
+                colors_by_ind[person_ind] = face_ind + 2
+            else:
+                colors_by_ind[face_ind] = 0
+        for person_ind in self.unassigned_persons_inds:
+            colors_by_ind[person_ind] = 1
+
+        names = self.yolo_results.names
+        annotator = Annotator(
+            deepcopy(self.yolo_results.orig_img if img is None else img),
+            line_width,
+            font_size,
+            font,
+            pil,
+            example=names,
+        )
+        pred_boxes, show_boxes = self.yolo_results.boxes, boxes
+        pred_probs, show_probs = self.yolo_results.probs, probs
+
+        if pred_boxes and show_boxes:
+            for bb_ind, (d, age, gender, gender_score) in enumerate(
+                zip(pred_boxes, self.ages, self.genders, self.gender_scores)
+            ):
+                c, conf, guid = int(d.cls), float(
+                    d.conf) if conf else None, None if d.id is None else int(d.id.item())
+                name = ("" if guid is None else f"id:{guid} ") + names[c]
+                label = (
+                    f"{name} {conf:.2f}" if conf else name) if labels else None
+                if ages and age is not None:
+                    label += f" {age:.1f}"
+                if genders and gender is not None:
+                    label += f" {'F' if gender == 'female' else 'M'}"
+                if gender_probs and gender_score is not None:
+                    label += f" ({gender_score:.1f})"
+                annotator.box_label(d.xyxy.squeeze(), label,
+                                    color=colors(colors_by_ind[bb_ind], True))
+
+        if pred_probs is not None and show_probs:
+            text = f"{', '.join(f'{names[j] if names else j} {pred_probs.data[j]:.2f}' for j in pred_probs.top5)}, "
+            annotator.text((32, 32), text, txt_color=(
+                255, 255, 255))  # TODO: allow setting colors
+
+        return annotator.result()
+
+    def set_tracked_age_gender(self, tracked_objects: Dict[int, List[AGE_GENDER_TYPE]]):
+        """
+        Update age and gender for objects based on history from tracked_objects.
+        Args:
+            tracked_objects (dict[int, list[AGE_GENDER_TYPE]]): info about tracked objects by guid
+        """
+
+        for face_ind, person_ind in self.face_to_person_map.items():
+            pguid = self._get_id_by_ind(person_ind)
+            fguid = self._get_id_by_ind(face_ind)
+
+            if fguid == -1 and pguid == -1:
+                # YOLO might not assign ids for some objects in some cases:
+                # https://github.com/ultralytics/ultralytics/issues/3830
+                continue
+            age, gender = self._gather_tracking_result(
+                tracked_objects, fguid, pguid)
+            if age is None or gender is None:
+                continue
+            self.set_age(face_ind, age)
+            self.set_gender(face_ind, gender, 1.0)
+            if pguid != -1:
+                self.set_gender(person_ind, gender, 1.0)
+                self.set_age(person_ind, age)
+
+        for person_ind in self.unassigned_persons_inds:
+            pid = self._get_id_by_ind(person_ind)
+            if pid == -1:
+                continue
+            age, gender = self._gather_tracking_result(
+                tracked_objects, -1, pid)
+            if age is None or gender is None:
+                continue
+            self.set_gender(person_ind, gender, 1.0)
+            self.set_age(person_ind, age)
+
+    def _get_id_by_ind(self, ind: Optional[int] = None) -> int:
+        if ind is None:
+            return -1
+        obj_id = self.yolo_results.boxes[ind].id
+        if obj_id is None:
+            return -1
+        return obj_id.item()
+
+    def get_bbox_by_ind(self, ind: int, im_h: int = None, im_w: int = None) -> torch.tensor:
+        bb = self.yolo_results.boxes[ind].xyxy.squeeze().type(torch.int32)
+        if im_h is not None and im_w is not None:
+            bb[0] = torch.clamp(bb[0], min=0, max=im_w - 1)
+            bb[1] = torch.clamp(bb[1], min=0, max=im_h - 1)
+            bb[2] = torch.clamp(bb[2], min=0, max=im_w - 1)
+            bb[3] = torch.clamp(bb[3], min=0, max=im_h - 1)
+        return bb
+
+    def set_age(self, ind: Optional[int], age: float):
+        if ind is not None:
+            self.ages[ind] = age
+
+    def set_gender(self, ind: Optional[int], gender: str, gender_score: float):
+        if ind is not None:
+            self.genders[ind] = gender
+            self.gender_scores[ind] = gender_score
+
+    @staticmethod
+    def _gather_tracking_result(
+        tracked_objects: Dict[int, List[AGE_GENDER_TYPE]],
+        fguid: int = -1,
+        pguid: int = -1,
+        minimum_sample_size: int = 10,
+    ) -> AGE_GENDER_TYPE:
+
+        assert fguid != -1 or pguid != -1, "Incorrect tracking behaviour"
+
+        face_ages = [r[0] for r in tracked_objects[fguid] if r[0]
+                     is not None] if fguid in tracked_objects else []
+        face_genders = [r[1] for r in tracked_objects[fguid]
+                        if r[1] is not None] if fguid in tracked_objects else []
+        person_ages = [r[0] for r in tracked_objects[pguid]
+                       if r[0] is not None] if pguid in tracked_objects else []
+        person_genders = [r[1] for r in tracked_objects[pguid]
+                          if r[1] is not None] if pguid in tracked_objects else []
+
+        if not face_ages and not person_ages:  # both empty
+            return None, None
+
+        # You can play here with different aggregation strategies
+        # Face ages - predictions based on face or face + person, depends on history of object
+        # Person ages - predictions based on person or face + person, depends on history of object
+
+        if len(person_ages + face_ages) >= minimum_sample_size:
+            age = aggregate_votes_winsorized(person_ages + face_ages)
+        else:
+            face_age = np.mean(face_ages) if face_ages else None
+            person_age = np.mean(person_ages) if person_ages else None
+            if face_age is None:
+                face_age = person_age
+            if person_age is None:
+                person_age = face_age
+            age = (face_age + person_age) / 2.0
+
+        genders = face_genders + person_genders
+        assert len(genders) > 0
+        # take mode of genders
+        gender = max(set(genders), key=genders.count)
+
+        return age, gender
+
+    def get_results_for_tracking(self) -> Tuple[Dict[int, AGE_GENDER_TYPE], Dict[int, AGE_GENDER_TYPE]]:
+        """
+        Get objects from current frame
+        """
+        persons: Dict[int, AGE_GENDER_TYPE] = {}
+        faces: Dict[int, AGE_GENDER_TYPE] = {}
+
+        names = self.yolo_results.names
+        pred_boxes = self.yolo_results.boxes
+        for _, (det, age, gender, _) in enumerate(zip(pred_boxes, self.ages, self.genders, self.gender_scores)):
+            if det.id is None:
+                continue
+            cat_id, _, guid = int(det.cls), float(det.conf), int(det.id.item())
+            name = names[cat_id]
+            if name == "person":
+                persons[guid] = (age, gender)
+            elif name == "face":
+                faces[guid] = (age, gender)
+
+        return persons, faces
+
+    def associate_faces_with_persons(self):
+        face_bboxes_inds: List[int] = self.get_bboxes_inds("face")
+        person_bboxes_inds: List[int] = self.get_bboxes_inds("person")
+
+        face_bboxes: List[torch.tensor] = [
+            self.get_bbox_by_ind(ind) for ind in face_bboxes_inds]
+        person_bboxes: List[torch.tensor] = [
+            self.get_bbox_by_ind(ind) for ind in person_bboxes_inds]
+
+        self.face_to_person_map = {ind: None for ind in face_bboxes_inds}
+        assigned_faces, unassigned_persons_inds = assign_faces(
+            person_bboxes, face_bboxes)
+
+        for face_ind, person_ind in enumerate(assigned_faces):
+            face_ind = face_bboxes_inds[face_ind]
+            person_ind = person_bboxes_inds[person_ind] if person_ind is not None else None
+            self.face_to_person_map[face_ind] = person_ind
+
+        self.unassigned_persons_inds = [
+            person_bboxes_inds[person_ind] for person_ind in unassigned_persons_inds]
+
+    def crop_object(
+        self, full_image: np.ndarray, ind: int, cut_other_classes: Optional[List[str]] = None
+    ) -> Optional[np.ndarray]:
+
+        IOU_THRESH = 0.000001
+        MIN_PERSON_CROP_AFTERCUT_RATIO = 0.4
+        CROP_ROUND_RATE = 0.3
+        MIN_PERSON_SIZE = 50
+
+        obj_bbox = self.get_bbox_by_ind(ind, *full_image.shape[:2])
+        x1, y1, x2, y2 = obj_bbox
+        cur_cat = self.yolo_results.names[int(
+            self.yolo_results.boxes[ind].cls)]
+        # get crop of face or person
+        obj_image = full_image[y1:y2, x1:x2].copy()
+        crop_h, crop_w = obj_image.shape[:2]
+
+        if cur_cat == "person" and (crop_h < MIN_PERSON_SIZE or crop_w < MIN_PERSON_SIZE):
+            return None
+
+        if not cut_other_classes:
+            return obj_image
+
+        # calc iou between obj_bbox and other bboxes
+        other_bboxes: List[torch.tensor] = [
+            self.get_bbox_by_ind(other_ind, *full_image.shape[:2]) for other_ind in range(len(self.yolo_results.boxes))
+        ]
+
+        iou_matrix = box_iou(torch.stack([obj_bbox]), torch.stack(
+            other_bboxes)).cpu().numpy()[0]
+
+        # cut out other objects in case of intersection
+        for other_ind, (det, iou) in enumerate(zip(self.yolo_results.boxes, iou_matrix)):
+            other_cat = self.yolo_results.names[int(det.cls)]
+            if ind == other_ind or iou < IOU_THRESH or other_cat not in cut_other_classes:
+                continue
+            o_x1, o_y1, o_x2, o_y2 = det.xyxy.squeeze().type(torch.int32)
+
+            # remap current_person_bbox to reference_person_bbox coordinates
+            o_x1 = max(o_x1 - x1, 0)
+            o_y1 = max(o_y1 - y1, 0)
+            o_x2 = min(o_x2 - x1, crop_w)
+            o_y2 = min(o_y2 - y1, crop_h)
+
+            if other_cat != "face":
+                if (o_y1 / crop_h) < CROP_ROUND_RATE:
+                    o_y1 = 0
+                if ((crop_h - o_y2) / crop_h) < CROP_ROUND_RATE:
+                    o_y2 = crop_h
+                if (o_x1 / crop_w) < CROP_ROUND_RATE:
+                    o_x1 = 0
+                if ((crop_w - o_x2) / crop_w) < CROP_ROUND_RATE:
+                    o_x2 = crop_w
+
+            obj_image[o_y1:o_y2, o_x1:o_x2] = 0
+
+        obj_image, remain_ratio = cropout_black_parts(
+            obj_image, CROP_ROUND_RATE)
+        if remain_ratio < MIN_PERSON_CROP_AFTERCUT_RATIO:
+            return None
+
+        return obj_image
+
+    def collect_crops(self, image) -> PersonAndFaceCrops:
+
+        crops_data = PersonAndFaceCrops()
+        for face_ind, person_ind in self.face_to_person_map.items():
+            face_image = self.crop_object(
+                image, face_ind, cut_other_classes=[])
+
+            if person_ind is None:
+                crops_data.crops_faces_wo_body[face_ind] = face_image
+                continue
+
+            person_image = self.crop_object(
+                image, person_ind, cut_other_classes=["face", "person"])
+
+            crops_data.crops_faces[face_ind] = face_image
+            crops_data.crops_persons[person_ind] = person_image
+
+        for person_ind in self.unassigned_persons_inds:
+            person_image = self.crop_object(
+                image, person_ind, cut_other_classes=["face", "person"])
+            crops_data.crops_persons_wo_face[person_ind] = person_image
+
+        # uncomment to save preprocessed crops
+        # crops_data.save()
+        return crops_data

mivolo/version.py

@@ -0,0 +1 @@
+__version__ = "0.3.0dev"

requirements.txt

@@ -0,0 +1,6 @@
+ultralytics==8.0.187
+timm==0.8.13.dev0
+tqdm
+requests
+opencv-python
+omegaconf

utils.py

@@ -0,0 +1,60 @@
+import os
+import requests
+from tqdm import tqdm
+from modelscope import snapshot_download
+from urllib.parse import urlparse
+
+MODEL_DIR = snapshot_download("MuGeminorum/MiVOLO", cache_dir="./mivolo/__pycache__")
+
+
+def is_url(s: str):
+    try:
+        # 解析字符串
+        result = urlparse(s)
+        # 检查scheme（如http, https）和netloc（域名）
+        return all([result.scheme, result.netloc])
+
+    except:
+        # 如果解析过程中发生异常，则返回False
+        return False
+
+
+def download_file(url: str, save_path: str):
+    if os.path.exists(save_path):
+        print("目标已存在，无需下载")
+        return
+
+    create_dir(os.path.dirname(save_path))
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get("content-length", 0))
+    # 使用 tqdm 创建一个进度条
+    progress_bar = tqdm(total=total_size, unit="B", unit_scale=True)
+    with open(save_path, "wb") as file:
+        for data in response.iter_content(chunk_size=1024):
+            file.write(data)
+            progress_bar.update(len(data))
+
+    progress_bar.close()
+    if total_size != 0 and progress_bar.n != total_size:
+        os.remove(save_path)
+        print("下载失败，重试中...")
+        download_file(url, save_path)
+
+    else:
+        print("下载完成")
+
+    return save_path
+
+
+def create_dir(dir_path: str):
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+
+
+def get_jpg_files(folder_path: str):
+    all_files = os.listdir(folder_path)
+    return [
+        os.path.join(folder_path, file)
+        for file in all_files
+        if file.lower().endswith(".jpg")
+    ]