File size: 1,100 Bytes
b5f33fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
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("图片中没有检测到人体。")
|