import torch
import torchvision.transforms as transforms
from PIL import Image
from torchvision.models.detection import fasterrcnn_resnet50_fpn


def has_person(image_path):
    # 加载预训练的 Faster R-CNN 模型
    model = fasterrcnn_resnet50_fpn(pretrained=True)
    model.eval()

    # 载入并预处理图片
    img = Image.open(image_path)
    transform = transforms.Compose([transforms.ToTensor()])
    input_tensor = transform(img)
    input_batch = input_tensor.unsqueeze(0)

    # 模型推理
    with torch.no_grad():
        output = model(input_batch)

    # 解析输出结果
    labels = output[0]['labels'].numpy()
    scores = output[0]['scores'].numpy()

    # 判断是否检测到人体（label=1 表示人类类别）
    person_detected = any(label == 1 and score >
                          0.5 for label, score in zip(labels, scores))

    return person_detected


if __name__ == "__main__":
    image_path = './images/test.jpg'
    if has_person(image_path):
        print("图片中检测到人体。")
    else:
        print("图片中没有检测到人体。")