环境搭建
pipinstallopencv-python numpy
目录结构
chkimg.py models/ ├── efficientdet_d2/# EfficientDet-D2 模型文件│ ├── saved_model.pb# 模型文件│ ├── variables/# 模型变量│ └── assets/# 模型资源├── labels/# 标签文件│ └── coco_labels.txt# COCO数据集标签└── README.md# 本文件
模型下载说明
- EfficientDet-D2 模型
- 从 TensorFlow Hub 下载:https://tfhub.dev/tensorflow/efficientdet/d2/1
- 下载完成后,解压到
models/efficientdet_d2/目录。
- COCO 标签文件
- 从 GitHub 下载:https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/data/mscoco_label_map.pbtxt
- 下载完成后,保存为
models/labels/coco_labels.txt。
代码
#!/usr/bin/env python3# -*- coding: utf-8 -*-""" 物体检测脚本 - 完全本地化版本 使用OpenCV的DNN模块和预训练模型,无需TensorFlow Hub """importosimportsysimportcv2importnumpyasnpimportargparseclassLocalObjectDetector:def__init__(self):self.load_models()defload_models(self):"""加载本地物体检测模型"""try:print("正在加载本地物体检测模型...")# 检查本地模型文件model_dir=os.path.join(os.path.dirname(__file__),"models")dnn_model_path=os.path.join(model_dir,"yolov3.weights")dnn_config_path=os.path.join(model_dir,"yolov3.cfg")labels_path=os.path.join(model_dir,"labels","coco_labels.txt")# 如果本地模型不存在,使用OpenCV内置的DNN模型ifos.path.exists(dnn_model_path)andos.path.exists(dnn_config_path):print("使用本地YOLO模型...")self.net=cv2.dnn.readNet(dnn_model_path,dnn_config_path)self.use_yolo=Trueelse:print("本地模型不存在,使用OpenCV内置模型...")print("提示:您可以下载YOLO模型到 models/ 目录以获得更好的检测效果")self.use_yolo=False# 使用OpenCV的HOG行人检测器作为备选self.hog=cv2.HOGDescriptor()self.hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())# 加载COCO标签ifos.path.exists(labels_path):print("使用本地标签文件...")self.labels=self.load_coco_labels(labels_path)else:print("使用内置标签...")self.labels=self.get_builtin_labels()print("模型加载完成!")exceptExceptionase:print(f"模型加载失败:{e}")print("请确保已安装OpenCV")sys.exit(1)defload_coco_labels(self,labels_path):"""从文件加载COCO标签"""labels={}try:withopen(labels_path,'r',encoding='utf-8')asf:forlineinf:if'id:'inlineand'name:'inline:# 解析标签文件parts=line.strip().split()iflen(parts)>=4:id_part=parts[1]name_part=parts[3].strip('"')try:label_id=int(id_part)labels[label_id]=name_partexceptValueError:continueexceptExceptionase:print(f"加载标签文件失败:{e}")returnself.get_builtin_labels()returnlabelsdefget_builtin_labels(self):"""获取内置标签"""return{0:'person',1:'bicycle',2:'car',3:'motorcycle',4:'airplane',5:'bus',6:'train',7:'truck',8:'boat',9:'traffic light',10:'fire hydrant',11:'stop sign',12:'parking meter',13:'bench',14:'bird',15:'cat',16:'dog',17:'horse',18:'sheep',19:'cow',20:'elephant',21:'bear',22:'zebra',23:'giraffe',24:'backpack',25:'umbrella',26:'handbag',27:'tie',28:'suitcase',29:'frisbee',30:'skis',31:'snowboard',32:'sports ball',33:'kite',34:'baseball bat',35:'baseball glove',36:'skateboard',37:'surfboard',38:'tennis racket',39:'bottle',40:'wine glass',41:'cup',42:'fork',43:'knife',44:'spoon',45:'bowl',46:'banana',47:'apple',48:'sandwich',49:'orange',50:'broccoli',51:'carrot',52:'hot dog',53:'pizza',54:'donut',55:'cake',56:'chair',57:'couch',58:'potted plant',59:'bed',60:'dining table',61:'toilet',62:'tv',63:'laptop',64:'mouse',65:'remote',66:'keyboard',67:'cell phone',68:'microwave',69:'oven',70:'toaster',71:'sink',72:'refrigerator',73:'book',74:'clock',75:'vase',76:'scissors',77:'teddy bear',78:'hair drier',79:'toothbrush'}defdetect_objects_yolo(self,img):"""使用YOLO模型检测物体"""height,width=img.shape[:2]# 预处理图像blob=cv2.dnn.blobFromImage(img,1/255.0,(416,416),swapRB=True,crop=False)self.net.setInput(blob)# 前向传播layer_names=self.net.getLayerNames()output_layers=[layer_names[i-1]foriinself.net.getUnconnectedOutLayers()]outputs=self.net.forward(output_layers)# 解析检测结果detections=[]foroutputinoutputs:fordetectioninoutput:scores=detection[5:]class_id=np.argmax(scores)confidence=scores[class_id]ifconfidence>0.5:# 置信度阈值center_x=int(detection[0]*width)center_y=int(detection[1]*height)w=int(detection[2]*width)h=int(detection[3]*height)x=int(center_x-w/2)y=int(center_y-h/2)detection_info={'class_id':class_id,'class_name':self.labels.get(class_id,f'Unknown-{class_id}'),'confidence':float(confidence),'bbox':[x,y,x+w,y+h]}detections.append(detection_info)returndetectionsdefdetect_objects_hog(self,img):"""使用HOG检测器检测行人"""# 检测行人boxes,weights=self.hog.detectMultiScale(img,winStride=(8,8),padding=(4,4),scale=1.05)detections=[]for(x,y,w,h),weightinzip(boxes,weights):detection_info={'class_id':0,'class_name':'person','confidence':float(weight),'bbox':[x,y,x+w,y+h]}detections.append(detection_info)returndetectionsdefdetect_objects(self,image_path):"""检测图片中的物体"""try:# 读取图像img=cv2.imread(image_path)ifimgisNone:print(f"无法读取图像:{image_path}")returnNone,[]print("正在进行物体检测...")# 根据可用模型选择检测方法ifself.use_yolo:detections=self.detect_objects_yolo(img)else:detections=self.detect_objects_hog(img)returnimg,detectionsexceptExceptionase:print(f"物体检测失败:{e}")returnNone,[]defdraw_detections(self,img,detections):"""在图像上绘制检测结果"""ifimgisNoneornotdetections:returnimgfordetectionindetections:bbox=detection['bbox']class_name=detection['class_name']confidence=detection['confidence']# 绘制边界框x1,y1,x2,y2=bbox cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),2)# 绘制标签label=f"{class_name}:{confidence:.2f}"label_size=cv2.getTextSize(label,cv2.FONT_HERSHEY_SIMPLEX,0.5,2)[0]cv2.rectangle(img,(x1,y1-label_size[1]-10),(x1+label_size[0],y1),(0,255,0),-1)cv2.putText(img,label,(x1,y1-5),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),2)returnimgdefanalyze_image(self,image_path):"""分析图像中的物体"""ifnotos.path.exists(image_path):print(f"错误: 文件不存在{image_path}")returnprint(f"\n{'='*50}")print(f"分析图像:{image_path}")print(f"{'='*50}")# 基本信息try:img=cv2.imread(image_path)ifimgisnotNone:height,width,channels=img.shapeprint(f"图像尺寸:{width}x{height}")print(f"颜色通道:{channels}")file_size=os.path.getsize(image_path)/1024print(f"文件大小:{file_size:.2f}KB")else:print("无法读取图像")returnexceptExceptionase:print(f"读取图像失败:{e}")return# 物体检测print(f"\n{'='*20}物体检测{'='*20}")result_img,detections=self.detect_objects(image_path)ifdetections:print(f"检测到{len(detections)}个物体:")print("-"*40)# 按置信度排序detections.sort(key=lambdax:x['confidence'],reverse=True)fori,detectioninenumerate(detections,1):print(f"{i}.{detection['class_name']}")print(f" 置信度:{detection['confidence']:.3f}")print(f" 位置:{detection['bbox']}")print()# 统计物体类型object_counts={}fordetectionindetections:class_name=detection['class_name']object_counts[class_name]=object_counts.get(class_name,0)+1print("物体统计:")forobj_type,countinobject_counts.items():print(f"{obj_type}:{count}个")else:print("未检测到任何物体")# 保存结果ifresult_imgisnotNoneanddetections:result_img=self.draw_detections(result_img,detections)output_path=f"detected_{os.path.basename(image_path)}"cv2.imwrite(output_path,result_img)print(f"\n检测结果已保存:{output_path}")defbatch_process(self,folder_path):"""批量处理文件夹中的图像"""ifnotos.path.isdir(folder_path):print(f"错误:{folder_path}不是有效文件夹")returnextensions=['.jpg','.jpeg','.png','.bmp','.tiff']image_files=[]forextinextensions:image_files.extend([fforfinos.listdir(folder_path)iff.lower().endswith(ext)])ifnotimage_files:print(f"在{folder_path}中未找到图像文件")returnprint(f"找到{len(image_files)}个图像文件")fori,image_fileinenumerate(image_files,1):image_path=os.path.join(folder_path,image_file)print(f"\n处理{i}/{len(image_files)}:{image_file}")self.analyze_image(image_path)defmain():parser=argparse.ArgumentParser(description='本地物体检测工具')parser.add_argument('image_path',nargs='?',help='要检测的图片路径')parser.add_argument('--batch',help='批量处理文件夹')parser.add_argument('--confidence',type=float,default=0.5,help='检测置信度阈值 (默认: 0.5)')args=parser.parse_args()ifnotargs.image_pathandnotargs.batch:print("请提供图像路径或使用 --batch 参数指定文件夹")parser.print_help()returndetector=LocalObjectDetector()ifargs.batch:detector.batch_process(args.batch)else:detector.analyze_image(args.image_path)if__name__=="__main__":main()
代码运行命令
python chkimg.py 图片路径.jpg python chkimg.py--batch文件夹路径 python chkimg.py-h
识别结果
![]()
![]()
![]()
![]()
![]()