human-detector / human_detect.py
George
upl all codes
b5f33fd
raw
history blame
No virus
1.1 kB
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("图片中没有检测到人体。")