from data.dataloader_detection import ListDataset
from data.collate_function import collate_function
from config.config import load_config, cfg
from net.resnet import ResNet, ResnetBasic,ResnetBasicSlim
from utils.nms_np_simple import nms
from data.dataloader_test_detection import ListDataset as testListDataset
from data.resize_uniform import resizeUniform
from dataY.yolov1_dataY import DataY
from net.yolov1 import YOLOv1
import numpy as np
import cv2
from loss.yololoss import yoloLoss
from dataY.yolov1_dataY import DataY
import time
import torch
def plotGride(img, grideHW=(7,7), stride=64):
# plot gride
h,w,_ = img.shape
for li in range(grideHW[1]):
cv2.line(img, (int(li *stride), 0),
(int(li * stride), int(w)), (0, 255, 0), 1)
for li in range(grideHW[0]):
cv2.line(img, (0, int(li * stride)),
(int(h), int(li * stride)), (0, 255, 0), 1)
return img
def plotBox(img, x1, y1,x2,y2,txts):
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 1)
for i in range(len(txts)):
cv2.putText(img, txts[i], (int(x1), int(y1 + i*15)), 1, 1, (0, 0, 255))
cv2.circle(img, (int((x1 +x2) / 2), int((y1 + y2) / 2)), color=(0, 0, 255), radius=2, thickness=-1)
return img
def post(pred, dic):
# target contain object mask,
b, h, w, c = pred.size()
alldets= []# [[一个图片的dets]，[],...]
for i in range(b):
apred = pred[i]
"""根据置信度来判断哪个网络有"""
coobjMask1 = apred[:,:,4] > dic["scoreThresh"]
for j in range(dic["bboxPredNum"]):
coobjMask1 = torch.logical_or(coobjMask1, apred[:, :, 4 + i * 5] > dic["scoreThresh"])
"""每个cell 选择最大置信度的哪个"""
confPred = apred[:, :, 4:5 * dic["bboxPredNum"]:5]
_, maxIndex = torch.max(confPred, dim=-1, keepdim=True)
h_, w_, c_ = apred.size()
oneHot = torch.zeros(h_, w_, dic["bboxPredNum"]).to("cuda:0").scatter_(-1, maxIndex, 1).type(
torch.bool)
"""最终选择"""
coobjMask = torch.logical_and(oneHot, torch.unsqueeze(coobjMask1, -1))
chioceIndex = torch.nonzero(coobjMask, as_tuple=False)
chioceIndex = chioceIndex.to('cpu').numpy()
dets = []
for k in range(chioceIndex.shape[0]):
hid, wid, boxid = chioceIndex[k][0], chioceIndex[k][1], chioceIndex[k][2]
fullVector = apred[hid][wid].to("cpu").numpy()
clsVec = fullVector[-dic["clsNum"]:]
boxVector = fullVector[boxid * 5: (boxid + 1) * 5]
deltax, deltay, w, h, score = boxVector[0], boxVector[1], boxVector[2], boxVector[3], boxVector[4]
cy = (hid + deltay) * dic["stride"]
cx = (wid + deltax) * dic["stride"]
w = w * w * dic["netInputHw"][1]
h = h * h * dic["netInputHw"][0]
c = clsVec.argmax()
dets.append([cx, cy, w, h, score, c])
dets = np.array(dets)
dets = dets.reshape(-1, 6)
dets[:, :2] -= dets[:, 2:4] / 2
dets[:, 2:4] += dets[:, :2]
#boarder in the image
dets[:,:2] = np.where(dets[:, :2] < 0, 0, dets[:,:2] )
dets[:, 2] = np.where(dets[:, 2] > dic["netInputHw"][1], dic["netInputHw"][1], dets[:, 2])
dets[:, 3] = np.where(dets[:, 3] > dic["netInputHw"][0], dic["netInputHw"][0], dets[:, 3])
dets = nms(dets, dic["iouThresh"])
alldets.append(dets)
return alldets
if __name__ == '__main__':
"""config"""
load_config(cfg, "./config/config.yaml")
print(cfg)
device = torch.device('cuda:0')
batchsize = 2
showFlag = 1
saveFlag = 0
saveDir = ""
mode = 1
# 如果是test mode 没有label 信息， 只能是预测
# cam Mode 调用摄像头
# val model这个有label信息，在show的时候会展示label
modeDict = {0:"testMode", 1:"valMode", 2:"camMode"}
postDict = {"scoreThresh": 0.3,
"iouThresh": 0.3,
"netInputHw":(448,448),
"bboxPredNum": cfg.model.bboxPredNum,
"clsNum":cfg.model.clsNum,
"stride":cfg.model.stride,
}
"""dataset"""
if modeDict[mode] == "testMode":
dataset = testListDataset( imgPath = cfg.dir.testImgDir, # images root /
netInputSizehw = cfg.model.netInput,
imgChannelNumber=cfg.model.imgChannelNumber,
clsname= cfg.clsname
)
if modeDict[mode] == "valMode":
dataset = ListDataset(trainAnnoPath=cfg.dir.valAnnoDir, trainImgPath=cfg.dir.valImgDir,
netInputSizehw=cfg.model.netInput, augFlag=False,
clsname = cfg.clsname, imgChannelNumber=cfg.model.imgChannelNumber)
if modeDict[mode] == "valMode" or modeDict[mode] == "testMode":
dataLoader = torch.utils.data.DataLoader(
dataset,
collate_fn=collate_function,
batch_size= batchsize, #cfg.train.batchSize,
shuffle= False,
num_workers=cfg.train.workers,
pin_memory=True, # 如果机器计算能力好的话，就可以设置为True，
)
datay = DataY(inputHW=cfg.model.netInput, # 指定了inputsize 这是因为输入的是经过resize后的图片
gride=cfg.model.featSize, # 将网络输入成了多少个网格
stride=cfg.model.stride,
boxNum=cfg.model.bboxPredNum,
clsNum=cfg.model.clsNum)
"""准备网络"""
# network = ResNet(ResnetBasic, [2, 2, 2, 2], channel_out = 15)
network = ResNet(ResnetBasicSlim,
[2, 2, 2, 2],
# [3,4,6,3],
channel_in=cfg.data.imgChannelNumber,
channel_out=(cfg.model.bboxPredNum * 5 + cfg.model.clsNum))
# network = YOLOv1(params={"dropout": 0.5, "num_class": cfg.model.clsNum})
network.to(device)
weights = torch.load(cfg.dir.modelSaveDir + cfg.dir.modelName) # 加载参数
network.load_state_dict(weights) # 给自己的模型加载参数
with torch.no_grad():
if modeDict[mode] == "camMode":
cap = cv2.VideoCapture(0)
while (1):
assert cfg.model.imgChannelNumber == 3, "输入通道目前支支持3通道"
ret, img = cap.read()
frame = np.copy(img)
if not ret:
print("can not cap a picture")
time.sleep(1)
continue
img, infos = resizeUniform(img, cfg.model.netInput)
imgs = np.array([img])
imgs = torch.from_numpy(imgs.astype(np.float32)),
imgs = imgs.to(device).float()
mean = torch.tensor(cfg.data.normalize[0]).cuda().reshape(3, 1, 1)
std = torch.tensor(cfg.data.normalize[1]).cuda().reshape(3, 1, 1)
imgs = (imgs - mean) / std
pred = network(imgs)
pred = pred.permute(0, 2, 3, 1) # rotate the BCHW to BHWC
"""post"""
bcDets = post(pred, postDict)
dets = bcDets[0]
cv2.imshow("capture", img)
"""plot pred"""
imgp = plotGride(frame, grideHW=(cfg.model.featSize), stride=cfg.model.stride)
for i in range(dets.shape[0]):
x1, y1, x2, y2, score, cls = \
dets[i][0], dets[i][1], dets[i][2], dets[i][3], dets[i][4], dets[i][5]
imgp = plotBox(imgp, x1, y1, x2, y2,
["s: " + str(round(score, 3)), "c: " + cfg.clsname[cls]])
cv2.imshow("pred", imgp)
cv2.waitKey()
if cv2.waitKey(1) & 0xFF == 32:
break
if modeDict[mode] == "valMode" or modeDict[mode] == "testMode":
for id, infos in enumerate(dataLoader): #每个batch
"""forward and pred"""
imgs = infos['images']
imgs = imgs.to(device).float()
mean = torch.tensor(cfg.data.normalize[0]).cuda().reshape(3, 1, 1)
std = torch.tensor(cfg.data.normalize[1]).cuda().reshape(3, 1, 1)
imgs = (imgs - mean) / std
pred = network(imgs)
pred = pred.permute(0, 2, 3, 1) # rotate the BCHW to BHWC
"""post"""
bcDets = post(pred, postDict)
for bcid in range(batchsize):
dets = bcDets[bcid]
if showFlag:
image = infos['images'][bcid]
image = image.to('cpu').numpy()
image = image.transpose(1, 2, 0).astype(np.uint8)
image = cv2.UMat(image).get()
imgt = np.copy(image)
imgp = np.copy(image)
"""plot pred"""
imgp = plotGride(imgp, grideHW=(cfg.model.featSize), stride=cfg.model.stride)
for i in range(dets.shape[0]):
x1, y1, x2, y2, score, cls = \
dets[i][0], dets[i][1], dets[i][2], dets[i][3], dets[i][4], dets[i][5]
imgp = plotBox(imgp, x1, y1, x2, y2,
[str(i)+" s: " + str(round(score, 3)), "c: " + cfg.clsname[cls]])
print("pred[%d x1:%.2f y1:%.2f w:%.2f h:%.2f score: %.2f "%(i, x1,y1,x2-x1,y2-y1,score)+ cfg.clsname[cls] +"]")
cv2.imshow("pred", imgp)
cv2.waitKey()
if showFlag and not modeDict[mode] == "testModeFlag":# test mode 没有target
"""read target"""
bboxesGt = infos['bboxesGt'][bcid]
classesGt = infos['classes'][bcid]
annoName = infos["annoName"][bcid]
"""plot target"""
imgt = plotGride(imgt, grideHW=(cfg.model.featSize), stride=cfg.model.stride)
for i in range(bboxesGt.shape[0]):
x1, y1, w, h = bboxesGt[i]
cls = classesGt[i]
plotBox(imgt, x1,y1,x1+w,y1+h,[cfg.clsname[cls]])
print("target[x1:%.2f y1:%.2f w:%.2f h:%.2f " % (x1, y1, w, h)+ cfg.clsname[cls]+"]")
print(annoName)
cv2.imshow("target", imgt)
cv2.waitKey()
print("-"*50)
# """post"""
# # target contain object mask
# coobjMask_ = pred[:, :, :, 4] > scoreThresh
# for i in range(1, cfg.model.bboxPredNum):
# coobjMask_ = torch.logical_or(coobjMask_, pred[:, :, :, 4 + i * 5] > scoreThresh)
#
# #选择每个anchor 点 预测置信度最大的, bbox num = 1 代码也可以
# confPred = pred[:, :, :, 4:5 * cfg.model.bboxPredNum:5]
# _, maxIndex = torch.max(confPred, dim=-1, keepdim=True)
# b_, h_, w_, c_ = pred.size()
# oneHot = torch.zeros(b_, h_, w_, cfg.model.bboxPredNum).to("cuda:0").scatter_(-1, maxIndex, 1).type(
# torch.bool)
#
# coobjMask = torch.logical_and(oneHot, torch.unsqueeze(coobjMask_, -1))
# chioceIndex = torch.nonzero(coobjMask, as_tuple=False) #
#
# """ plot pred """
# chioceIndex = chioceIndex.to('cpu').numpy()
# dets = []
# imgp = np.copy(image)
# for i in range(chioceIndex.shape[0]):
# imId, hid, wid, boxid = chioceIndex[i][0], chioceIndex[i][1], chioceIndex[i][2], chioceIndex[i][3]
#
# fullVector = pred[imId][hid][wid].to("cpu").numpy()
# clsVec = fullVector[-cfg.model.clsNum:]
# boxVector = fullVector[boxid * 5: (boxid + 1) * 5]
# deltax, deltay, w, h, score = boxVector[0], boxVector[1], boxVector[2], boxVector[3], boxVector[4]
#
# cy = (hid + deltay) * cfg.model.stride
# cx = (wid + deltax) * cfg.model.stride
# w = w*w* cfg.model.netInput[1]
# h = h*h * cfg.model.netInput[0]
# c = clsVec.argmax()
# dets.append([cx, cy, w, h, score, c])
# dets = np.array(dets)
# dets = dets.reshape(-1, 6)
# dets[:, :2] -= dets[:, 2:4] / 2
# dets[:, 2:4] += dets[:, :2]
# # dets = nms(dets, iouThresh)
if saveFlag and modeDict[mode] == "valModeFlag":
np.savetxt(saveDir + infos["annoName"][bcid],dets)
if saveFlag and modeDict[mode] == "testModeFlag":
np.savetxt(saveDir + infos["imgName"][bcid].split["."][0]+".txt",dets)
