def cam_inference(self, fname):
cap, size, _, fps = self._video_params(0)
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
if fname: writer = cv2.VideoWriter(fname, fourcc, fps, size)
while (cap.isOpened()):
ret, frame = cap.read()
if ret == True:
boxes, scores, _, labels = self.inf_model.predict(frame)
frame_pred = draw_boxes(frame, (boxes, scores, labels))
if fname: writer.write(frame_pred)
cv2.imshow('Yolo Output', frame_pred)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
if fname:
writer.release()
if YoloParams.STORE_GIF: generate_gif(fname)
cv2.destroyAllWindows()
def inference(self, path):
flag = self.debug_timings
if os.path.isdir(path):
fnames = [
os.path.join(path, f) for f in os.listdir(path)
if os.path.isfile(os.path.join(path, f))
]
out_fname_mod = '.png'
out_path = os.path.join(path, 'out')
mkdir_p(out_path)
else:
fnames = [path]
out_fname_mod = '_pred.png'
out_path = os.path.dirname(path)
flag = False
for f in tqdm(fnames, desc='Processing Batch'):
image = cv2.imread(f)
plt.figure(figsize=(10, 10))
boxes, scores, _, labels = self.inf_model.predict(image.copy())
print(f, list(zip(labels, scores)))
image = draw_boxes(image, (boxes, scores, labels))
out_name = os.path.join(
out_path,
os.path.basename(f).split('.')[0] + out_fname_mod)
cv2.imwrite(out_name, image)
def video_inference(self, filename):
cap, size, video_len, fps = self._video_params(filename)
outname = filename.split('.')[0] + "_pred.mp4"
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
writer = cv2.VideoWriter(outname, fourcc, fps, size)
for i in tqdm(range(video_len)):
ret, frame = cap.read()
boxes, scores, _, labels = self.inf_model.predict(frame)
frame_pred = draw_boxes(frame, (boxes, scores, labels))
writer.write(frame_pred)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
writer.release()
cv2.destroyAllWindows()
if YoloParams.STORE_GIF: generate_gif(outname)