def detect_vedio(image, network, thresh, names):
transform = T.Compose([T.ToTensor(),T.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])
evaluator = eva.evalYolov2(network.layers[-1].flow[0], class_thresh=thresh)
if image == '0':
cap = cv2.VideoCapture(0)
else:
cap = cv2.VideoCapture(image)
if not cap.isOpened():
print("Unable to open camera")
exit(-1)
fps = 0.0
while(cap.isOpened()):
t0 = time.time()
ret, img_raw = cap.read()
b,g,r=cv2.split(img_raw)
img_raw1=cv2.merge([r,g,b])
h_im , w_im, _ = img_raw.shape
img = cv2.resize(img_raw1,(network.width, network.height))
if ret == True:
img = transform(img).view(1,3,network.height,network.width).cuda()
pred = network(img)
result = evaluator.forward(pred, w_im=w_im, h_im=h_im)
t1 = time.time()
fps = 1/(t1-t0)
if result is not None:
im = plot_boxes_cv2(img_raw, result, names, fps=fps)
else:
im = img_raw
cv2.imshow('prediction',im)
print('fps: %f'%fps)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
return
def detect_image(image, network, thresh, names):
pil_img = Image.open(image)
w_im, h_im = pil_img.size
transform = T.Compose([T.ToTensor(),T.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])
img = pil_img.resize( (network.width, network.height) )
img = transform(img).cuda()
img = img.view(1,3,network.height,network.width)
pred = network(img)
evaluator = eva.evalYolov2(network.layers[-1].flow[0], class_thresh=thresh)
result = evaluator.forward(pred, w_im=w_im, h_im=h_im)
print(result)
image1 = cv2.imread(image)
im = plot_boxes_cv2(image1, result, names)
cv2.imshow('prediction',im)
cv2.waitKey(0)
cv2.destroyAllWindows()
return
netname = args.netcfg.split('.')[0].split('/')[-1]
layer = []
print('the depth of the network is %d'%(layerList.__len__()-1))
network = net.network(layerList)
#step 2: load network parameters
network.load_weights(args.weight)
layerNum = network.layerNum
if args.cuda:
network = network.cuda()
#step 3: load data and test
network = network.eval()
if args.eval == 1:
dataset = dat.YoloDataset(validlist,416,416,train=0)
dataloader = data.DataLoader(dataset, batch_size=16, shuffle=1, drop_last=True)
dataiter = iter(dataloader)
evaluator = eva.evalYolov2(network.layers[-1].flow[0], class_thresh=args.thresh)
for i in range(100):
imgs, labels = next(dataiter)
if args.cuda == 1:
imgs = imgs.cuda()
labels = labels.cuda()
pred = network(imgs)
result = evaluator.forward(pred, truth=labels)
else:
image = args.img
img_tail = image.split('.')[-1]
if img_tail == 'jpg' or img_tail =='jpeg' or img_tail == 'png':
detect_image(image, network, args.thresh, names)
elif img_tail == 'mp4' or img_tail =='mkv' or img_tail == 'avi' or img_tail =='0':
detect_vedio( image, network, args.thresh, names)