151 lines
5.3 KiB
Python
Executable File
151 lines
5.3 KiB
Python
Executable File
'''
|
||
author ： Hatcher W
|
||
Email： hcwang@semi.ac.cn
|
||
use: 目标检测Y5推理
|
||
input:需要推理的文件目录
|
||
output：推理结果
|
||
'''
|
||
import time
|
||
import cv2
|
||
import torch,os
|
||
import numpy as np
|
||
from numpy import random
|
||
import os,sys
|
||
root_path = os.getcwd()
|
||
sys.path.insert(0,root_path+"/Y5")
|
||
from Y5.models.experimental import attempt_load
|
||
from Y5.utils.datasets import letterbox
|
||
from Y5.utils.general import non_max_suppression, \
|
||
scale_coords
|
||
from Y5.utils.plots import plot_one_box
|
||
|
||
|
||
|
||
class Detector:
|
||
|
||
def __init__(self, device):
|
||
super(Detector, self).__init__()
|
||
self.img_size = 640
|
||
self.conf_thres = 0.3
|
||
self.iou_thres = 0.4
|
||
self.stride = 1
|
||
|
||
self.weights = '../data_MPKFNet/weigth/detection/Y/best.pt'
|
||
self.device = device
|
||
#self.device = select_device(self.device)
|
||
model = attempt_load(self.weights, map_location=self.device)
|
||
model.to(self.device).eval()
|
||
model.half()
|
||
self.model = model
|
||
self.names = model.module.names if hasattr(
|
||
model, 'module') else model.names
|
||
|
||
def preprocess(self, img):
|
||
|
||
img0 = img.copy()
|
||
img = letterbox(img, new_shape=self.img_size, stride=self.stride)[0]
|
||
img = img[:, :, ::-1].transpose(2, 0, 1)
|
||
img = np.ascontiguousarray(img)
|
||
img = torch.from_numpy(img).to(self.device)
|
||
img = img.half() # °ë¾«¶È
|
||
img /= 255.0 # ͼÏñ¹éÒ»»¯
|
||
if img.ndimension() == 3:
|
||
img = img.unsqueeze(0)
|
||
|
||
return img0, img
|
||
|
||
def detect(self, im):
|
||
|
||
im0, img = self.preprocess(im)
|
||
|
||
pred = self.model(img, augment=False)[0]
|
||
pred = pred.float()
|
||
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
|
||
|
||
pred_boxes = []
|
||
scores = []
|
||
labels = []
|
||
for det in pred:
|
||
|
||
if det is not None and len(det):
|
||
det[:, :4] = scale_coords(
|
||
img.shape[2:], det[:, :4], im0.shape).round()
|
||
|
||
pred_boxes.append(det[:,:4])
|
||
scores.append(det[:,4])
|
||
labels.append((det[:,-1]))
|
||
pred_boxes = pred_boxes[0]
|
||
scores = scores[0]
|
||
labels = labels[0]
|
||
# for *x, conf, cls_id in det:
|
||
# lbl = self.names[int(cls_id)]
|
||
# if not lbl in ['person', 'car', 'truck']:
|
||
# continue
|
||
# x1, y1 = int(x[0]), int(x[1])
|
||
# x2, y2 = int(x[2]), int(x[3])
|
||
# pred_boxes.append(
|
||
# (x1, y1, x2, y2, lbl, conf))
|
||
|
||
return im, pred_boxes,scores,labels
|
||
|
||
|
||
if __name__ == '__main__':
|
||
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '4'
|
||
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '4'
|
||
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||
#path = '/data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit/VOC5/JPEGImages/'
|
||
path = "/data/Hatcher/ouput1/0001_20210414151147/RawImage/"
|
||
outpath = "/data/Hatcher/detectout3"
|
||
if not os.path.exists(outpath):
|
||
os.makedirs(outpath)
|
||
detector = Detector(device)
|
||
path_lists = os.listdir(path)
|
||
for image_path in path_lists:
|
||
img = cv2.imread(os.path.join(path,image_path)) # BGR
|
||
|
||
# Get names and colors
|
||
names = [ 'background','aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
|
||
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]#detector.model.module.names if hasattr(detector.model, 'module') else detector.model.names
|
||
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
|
||
t3 = time.time()
|
||
im0, pred_boxes ,scores,labels = detector.detect(img)
|
||
t4 = time.time()
|
||
print("推理： ", (t4 - t3) * 1000)
|
||
print(type(labels))
|
||
if len(labels) == 0:
|
||
continue
|
||
indices_label_left = labels < 4
|
||
indices_label_right = labels > 2
|
||
indices = []
|
||
for i in range(len(indices_label_left)):
|
||
indices.append(indices_label_left[i] and indices_label_right[i])
|
||
|
||
pred_boxes = pred_boxes[indices]
|
||
labels = labels[indices]
|
||
scores = scores[indices] #因为莫名奇妙多了一个list维度，所以加了这个0
|
||
pred_boxes = pred_boxes[0]
|
||
labels = labels[0]
|
||
scores = scores[0]
|
||
print(len(scores))
|
||
for i in range(len(scores)):
|
||
xyxy = pred_boxes[i]
|
||
conf = scores[i]
|
||
cls = labels[i]
|
||
print(cls)
|
||
label = f'{names[int(cls)]} {conf:.2f}'
|
||
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
|
||
#drawn_image = draw_boxes(im0, pred_boxes, labels, scores, names).astype(np.uint8)
|
||
# for *xyxy, conf, cls in pred_boxes,scores,labels:
|
||
# label = f'{names[int(cls)]} {conf:.2f}'
|
||
# plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
|
||
|
||
# Stream results
|
||
#cv2.imshow(str(path), im0)
|
||
t1 = time.time()
|
||
cv2.imwrite(os.path.join(outpath, image_path),im0)
|
||
t2 = time.time()
|
||
print("写入： ",(t2-t1)*1000)
|
||
#cv2.waitKey(-1)
|