赞
踩
python烟火检测pytorch火焰抽烟正常检测EfficientNet
烟火检测功能主要应用于建筑工地、工业园区、仓库、以及其他易燃易爆场景的监控。该功能适用于白天或者光照条件较好的夜晚环境,不适用于光照条件较差、遮挡严重的场景。具体应用示例如下:
- #!/usr/bin/python3
- # -*- coding: utf-8 -*-
- '''
- '''
- import json
- from PIL import Image, ImageDraw, ImageFont
- import torch
- from torch import nn
- from torchvision import transforms
- from efficientnet_pytorch import FireSmokeEfficientNet
- import collections
- image_dir = './tests/5.jpg'
- model_para = collections.OrderedDict()
- model = FireSmokeEfficientNet.from_arch('efficientnet-b0')
- # out_channels = model._fc.in_features
- model._fc = nn.Linear(1280, 3)
- print(model)
- modelpara = torch.load('./checkpoint.pth.tar',map_location='cpu')
- # print(modelpara['state_dict'].keys())
- for key in modelpara['state_dict'].keys():
- # print(key[7:])
- # newkey = model_para[key.split('.',2)[-1]]
- # print(newkey)
- model_para[key[7:]] =modelpara['state_dict'][key]
-
- # print(model_para.keys())
- # 训练模型转换
-
-
-
- model.load_state_dict(model_para)
-
- # Preprocess image
- tfms = transforms.Compose([transforms.Resize(224), transforms.ToTensor(),
- transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
- image = Image.open(image_dir)
- img = tfms(image).unsqueeze(0)
- print(img.shape) # torch.Size([1, 3, 224, 224])
-
- # Load ImageNet class names
- labels_map = json.load(open('examples/simple/fire_smoke_map.txt'))
- labels_map = [labels_map[str(i)] for i in range(3)]
-
- # Classify
- model.eval()
- with torch.no_grad():
- outputs = model(img)
-
- draw = ImageDraw.Draw(image)
- font = ImageFont.truetype('simkai.ttf', 30)
- # Print predictions
- print('-----')
- cout = 0
- for idx in torch.topk(outputs, k=2).indices.squeeze(0).tolist():
- cout += 1
- prob = torch.softmax(outputs, dim=1)[0, idx].item()
- print('{label:<75} ({p:.2f}%)'.format(label=labels_map[idx], p=prob*100))
- position = (10, 30*cout - 20)
- text = '{label:<5} :{p:.2f}%'.format(label=labels_map[idx], p=prob*100)
- draw.text(position, text, font=font, fill="#ff0000", spacing=0, align='left')
-
- image.save('results/result_{}'.format(image_dir.split('/')[-1]))
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。