def write_video(vid_path, out_path, roi, endpoints, detections, scores):
"""Writes a new video with bounding boxes around detections.
# Arguments
vid_path: Path to the original video.
out_path: Path to the output video.
roi: Region of interest output from find_roi.
endpoints: Ruler endpoints from find_roi.
detections: Detections for each frame.
scores: Cover and species scores for each detection.
"""
# Initialize the video reader.
reader = openem.VideoReader()
status = reader.Init(vid_path)
if not status == openem.kSuccess:
raise IOError("Failed to read video {}!".format(vid_path))
# Initialize the video writer.
print("Writing annotated video to {}".format(out_path))
writer = openem.VideoWriter()
status = writer.Init(out_path, reader.FrameRate(), openem.kWmv2,
(reader.Width(), reader.Height()))
if not status == openem.kSuccess:
raise IOError("Failed to write video {}!".format(out_path))
# Iterate through frames.
for det_frame, score_frame in zip(detections, scores):
frame = openem.Image()
status = reader.GetFrame(frame)
if not status == openem.kSuccess:
raise RuntimeError("Error retrieving video frame!")
frame.DrawRect(roi, (255, 0, 0), 1, endpoints)
for j, (det, score) in enumerate(zip(det_frame, score_frame)):
clear = score.cover[2]
hand = score.cover[1]
if j == 0:
if clear > hand:
det_color = (0, 255, 0)
else:
det_color = (0, 0, 255)
frame.DrawRect(det.location, det_color, 2, endpoints, roi)
status = writer.AddFrame(frame)
if not status == openem.kSuccess:
raise RuntimeError("Error adding frame to video!")
def predict(config):
# Get paths from config file.
videos = config.test_vids()
out_dir = config.test_output_dir()
find_ruler_path = config.find_ruler_model_path()
detect_path = config.detect_model_path()
classify_path = config.classify_model_path()
count_path = config.count_model_path()
os.makedirs(out_dir, exist_ok=True)
# Iterate through videos.
for vid in videos:
# Get the video width.
reader = openem.VideoReader()
status = reader.Init(vid)
if not status == openem.kSuccess:
raise IOError("Failed to read video {}!".format(vid))
vid_width = float(reader.Width())
# Get path to output csv.
_, fname = os.path.split(vid)
base, _ = os.path.splitext(fname)
out_path = os.path.join(out_dir, base + '.csv')
# Find the ROI.
print("Doing prediction on {}...".format(vid))
roi, endpoints = _find_roi(find_ruler_path, vid)
# Do detection and classification.
detections, scores = _detect_and_classify(detect_path, classify_path,
vid, roi, endpoints)
# Write counts to csv.
_write_counts(count_path, out_path, roi, detections, scores, [
'unknown',
] + config.species(), vid_width, config.detect_width())
def _find_roi(mask_finder_path, vid_path):
"""Finds ROI in a video.
# Arguments
mask_finder_path: Path to find_ruler model file.
vid_path: Path to the video.
# Returns:
Region of interest and ruler endpoints.
# Raises:
IOError: If video or model file cannot be opened.
RuntimeError: If fails to add images or process model.
"""
# Determined by experimentation with GPU having 8GB memory.
max_img = 8
# Create and initialize the mask finder.
mask_finder = openem.RulerMaskFinder()
status = mask_finder.Init(mask_finder_path)
if not status == openem.kSuccess:
raise IOError("Failed to initialize mask finder!")
# Set up video reader.
reader = openem.VideoReader()
status = reader.Init(vid_path)
if not status == openem.kSuccess:
raise IOError("Failed to open video!")
# Decode the first 100 frames and take the average mask.
mask_avg = np.zeros((reader.Height(), reader.Width()))
num_masks = 0
masks = openem.VectorImage()
vid_end = False
for i in range(math.ceil(100 / max_img)):
for j in range(max_img):
img = openem.Image()
status = reader.GetFrame(img)
if not status == openem.kSuccess:
vid_end = True
break
status = mask_finder.AddImage(img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to add frame to mask finder!")
status = mask_finder.Process(masks)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process mask finder!")
for mask in masks:
mask.Resize(reader.Width(), reader.Height())
mask_data = mask.DataCopy()
mask_data = np.array(mask_data)
mask_data = np.reshape(mask_data,
(reader.Height(), reader.Width()))
mask_avg += mask_data
num_masks += 1
if vid_end:
break
# Convert mean mask from numpy to openem format.
mask_vec = mask_avg.copy()
mask_vec = mask_vec / np.max(mask_vec)
mask_vec = mask_vec * 255.0
mask_vec = mask_vec.reshape(-1).astype(np.uint8).tolist()
mask_img = openem.Image()
mask_img.FromData(mask_vec, reader.Width(), reader.Height(), 1)
# Now that we have the best mask, use this to compute the ROI.
endpoints = openem.RulerEndpoints(mask_img)
r_mask = openem.Rectify(mask_img, endpoints)
roi = openem.FindRoi(r_mask)
return (roi, endpoints)
def _detect_and_classify(detect_path, classify_path, vid_path, roi, endpoints):
"""Finds and classifies detections for all frames in a video.
# Arguments
detect_path: Path to detect model file.
classify_path: Path to classify model file.
vid_path: Path to the video.
roi: Region of interest output from find_roi.
endpoints: Ruler endpoints from find_roi.
# Returns
Detection rects and classification scores.
# Raises
IOError: If video or model files cannot be opened.
RuntimeError: If unable to add frame or process a model.
"""
# Determined by experimentation with GPU having 8GB memory.
max_img = 32
# Create and initialize the detector.
detector = openem.Detector()
status = detector.Init(detect_path, 0.5)
if not status == openem.kSuccess:
raise IOError("Failed to initialize detector!")
# Create and initialize the classifier.
classifier = openem.Classifier()
status = classifier.Init(classify_path, 0.5)
if not status == openem.kSuccess:
raise IOError("Failed to initialize classifier!")
# Initialize the video reader.
reader = openem.VideoReader()
status = reader.Init(vid_path)
if not status == openem.kSuccess:
raise IOError("Failed to open video {}!".format(vid_path))
# Iterate through frames.
vid_end = False
detections = []
scores = []
while True:
# Find detections.
dets = openem.VectorVectorDetection()
imgs = [openem.Image() for _ in range(max_img)]
for i, img in enumerate(imgs):
status = reader.GetFrame(img)
if not status == openem.kSuccess:
vid_end = True
break
img = openem.Rectify(img, endpoints)
img = openem.Crop(img, roi)
status = detector.AddImage(img)
imgs[i] = img
if not status == openem.kSuccess:
raise RuntimeError("Failed to add frame to detector!")
status = detector.Process(dets)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process detector!")
detections += dets
# Classify detections
for det_frame, img in zip(dets, imgs):
score_batch = openem.VectorClassification()
for det in det_frame:
det_img = openem.GetDetImage(img, det.location)
status = classifier.AddImage(det_img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to add frame to classifier!")
status = classifier.Process(score_batch)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process classifier!")
scores.append(score_batch)
if vid_end:
break
return (detections, scores)
type=str,
help="Path to pb file with classify model.")
parser.add_argument("count_model",
type=str,
help="Path to pb file with count model.")
parser.add_argument("video_paths",
type=str,
nargs="+",
help="One or more paths to video files.")
parser.add_argument("--no_video",
action="store_true",
help="Disable writing annotated video.")
args = parser.parse_args()
for i, video_path in enumerate(args.video_paths):
# Get the video width.
reader = openem.VideoReader()
status = reader.Init(video_path)
if not status == openem.kSuccess:
raise IOError("Failed to read video {}!".format(video_path))
vid_width = float(reader.Width())
# Find the ROI.
print("Finding region of interest...")
roi, endpoints = find_roi(args.find_ruler_model, video_path)
# Find detections and classify them.
print("Performing detection and classification...")
detections, scores, detect_width = detect_and_classify(
args.detect_model, args.classify_model, video_path, roi, endpoints)
# Write counts to csv.