def _extract_roi(paths):
img_path, roi_path = paths
img = openem.Image()
status = img.FromFile(img_path)
if status != openem.kSuccess:
print("Failed to read image {}".format(img_path))
else:
roi = openem.Rectify(img, ((x1, y1), (x2, y2)))
print("Saving ROI to: {}".format(roi_path))
roi.ToFile(roi_path)
def extract_dets(config):
"""Extracts detection images.
# Arguments:
config: ConfigInterface object.
"""
sys.path.append('../python')
import openem
# Create directories to store detections.
os.makedirs(config.train_dets_dir(), exist_ok=True)
# Open the detection results csv.
det_results = pd.read_csv(config.detect_inference_path())
bar = progressbar.ProgressBar(max_value=len(det_results),
redirect_stdout=True,
redirect_stderr=True)
threshold=0
if config.config.has_option('Detect', 'ExtractThreshold'):
threshold= config.config.getfloat('Detect', 'ExtractThreshold')
# Create the detection images.
for _, row in bar(det_results.iterrows()):
if row['det_conf'] < threshold:
continue
# Get the new path.
vid_id = row['video_id']
f = "{:04d}-conf{:04d}.jpg".format(row['frame'], row['det_conf'])
roi_path = os.path.join(config.train_rois_dir(), vid_id, f)
det_dir = os.path.join(config.train_dets_dir(), vid_id)
os.makedirs(det_dir, exist_ok=True)
det_path = os.path.join(det_dir, f)
# Extract detections.
print("Saving detection image to: {}".format(det_path))
roi = openem.Image()
roi.FromFile(roi_path)
rect = openem.Rect([row['x'], row['y'], row['w'], row['h']])
det = openem.GetDetImage(roi, rect)
det.ToFile(det_path)
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):
"""Runs classification model on extracted detections.
# Arguments
config: ConfigInterface object.
"""
# Import deployment library.
sys.path.append('../python')
import openem
# Make a dict to contain classification results.
class_data = {
'video_id': [],
'frame': [],
'no_fish': [],
'covered': [],
'clear': []
}
species_list = [
'_',
] + config.species()
for spc in species_list:
class_data['species_' + spc] = []
# Initialize classifier from deployment library.
classifier = openem.Classifier()
status = classifier.Init(config.classify_model_path())
if not status == openem.kSuccess:
raise IOError("Failed to initialize detector!")
w, h = classifier.ImageSize()
for img_path in config.train_dets():
# Get video id from path.
path, fname = os.path.split(img_path)
frame, _ = os.path.splitext(fname)
video_id = os.path.basename(os.path.normpath(path))
# Load in image.
img = openem.Image()
status = img.FromFile(img_path)
if not status == openem.kSuccess:
raise IOError("Failed to load image {}".format(img_path))
img.Resize(w, h)
# Add image to processing queue.
status = classifier.AddImage(img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to load image {}".format(img_path))
# Process loaded image.
scores = openem.VectorClassification()
status = classifier.Process(scores)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process image {}!".format(img_path))
# Write classification results.
for score in scores:
class_data['video_id'].append(video_id)
class_data['frame'].append(int(frame))
for spc, spc_score in zip(species_list, score.species):
class_data['species_' + spc].append(spc_score)
class_data['no_fish'].append(score.cover[0])
class_data['covered'].append(score.cover[1])
class_data['clear'].append(score.cover[2])
print("Finished classification on {}".format(img_path))
# Write classification results to csv.
os.makedirs(config.inference_dir(), exist_ok=True)
d = pd.DataFrame(class_data)
d.to_csv(config.classify_inference_path(), index=False)
help="Path to protobuf file containing model.",
type=str)
parser.add_argument("image_files",
help="Paths to one or more image files.",
nargs="+",
type=str)
args = parser.parse_args()
# Create and initialize classifier.
classifier = openem.Classifier()
status = classifier.Init(args.pb_file)
if not status == openem.kSuccess:
raise IOError("Failed to initialize classifier!")
# Load in images.
imgs = [openem.Image() for _ in args.image_files]
w, h = classifier.ImageSize()
for img, p in zip(imgs, args.image_files):
status = img.FromFile(p)
if not status == openem.kSuccess:
raise IOError("Failed to load image {}".format(p))
img.Resize(w, h)
# Add images to processing queue.
for img in imgs:
status = classifier.AddImage(img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to add image for processing!")
# Process the loaded images.
scores = openem.VectorClassification()
def predict(config):
"""Runs find ruler model on video frames.
# Arguments
config: ConfigInterface object.
"""
# Import deployment library.
sys.path.append('../python')
import openem
# Make a dict to contain find ruler results.
find_ruler_data = {
'video_id' : [],
'x1' : [],
'y1' : [],
'x2' : [],
'y2' : [],
}
# Make a dict to store mean of all masks.
mask_avg = {}
# Make a dict to store number of masks found per video.
num_masks = defaultdict(int)
# Create and initialize the mask finder.
mask_finder = openem.RulerMaskFinder()
status = mask_finder.Init(config.find_ruler_model_path())
if not status == openem.kSuccess:
raise IOError("Failed to initialize ruler mask finder!")
for img_path in config.train_imgs():
# Get video id from path.
path, fname = os.path.split(img_path)
frame, _ = os.path.splitext(fname)
video_id = os.path.basename(os.path.normpath(path))
if num_masks[video_id] > 200:
continue
else:
num_masks[video_id] += 1
print("Finding mask for image {}...".format(img_path))
# Load in image.
img = openem.Image()
status = img.FromFile(img_path)
if not status == openem.kSuccess:
continue
# Add image to processing queue.
status = mask_finder.AddImage(img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to add image {} for processing!".format(img_path))
# Process the loaded image.
masks = openem.VectorImage()
status = mask_finder.Process(masks)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process image {}!".format(img_path))
# Resize the mask back to the same size as the image.
mask = masks[0]
mask.Resize(img.Width(), img.Height())
# Initialize the mean mask if necessary.
if video_id not in mask_avg:
mask_avg[video_id] = np.zeros((img.Height(), img.Width()));
# Add the mask to the mask average.
mask_data = mask.DataCopy()
mask_data = np.array(mask_data)
mask_data = np.reshape(mask_data, (img.Height(), img.Width()))
mask_avg[video_id] += mask_data
for video_id in mask_avg:
print("Finding ruler endpoints for video {}...".format(video_id))
# Convert mask image from numpy to openem format.
mask_vec = mask_avg[video_id].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, img.Width(), img.Height(), 1);
# Get ruler endpoints from the mask averages.
p1, p2 = openem.RulerEndpoints(mask_img)
x1, y1 = p1
x2, y2 = p2
find_ruler_data['video_id'].append(video_id)
find_ruler_data['x1'].append(x1)
find_ruler_data['y1'].append(y1)
find_ruler_data['x2'].append(x2)
find_ruler_data['y2'].append(y2)
# Write detections to csv.
os.makedirs(config.inference_dir(), exist_ok=True)
d = pd.DataFrame(find_ruler_data)
d.to_csv(config.find_ruler_inference_path(), index=False)
def predict(config):
"""Runs detection model on extracted ROIs.
# Arguments
config: ConfigInterface object.
"""
# Import deployment library.
sys.path.append('../python')
import openem
# Make a dict to contain detection results.
det_data = {
'video_id' : [],
'frame' : [],
'x' : [],
'y' : [],
'w' : [],
'h' : [],
'det_conf' : [],
'det_species' : []
}
# Initialize detector from deployment library.
detector = openem.Detector()
status = detector.Init(config.detect_model_path())
if not status == openem.kSuccess:
raise IOError("Failed to initialize detector!")
limit = None
count = 0
threshold=0
if config.config.has_option('Detect', 'Limit'):
limit = config.config.getint('Detect','Limit')
if config.config.has_option('Detect', 'Threshold'):
threshold= config.config.getfloat('Detect', 'Threshold')
images=set()
for img_path in config.train_rois():
images.add(os.path.basename(os.path.dirname(img_path)))
if limit:
print(f"Limiting process to {limit} files.")
if len(images) >= limit:
break
else:
count = count + 1
# Load in image.
img = openem.Image()
status = img.FromFile(img_path)
if not status == openem.kSuccess:
raise IOError("Failed to load image {}".format(p))
# Add image to processing queue.
status = detector.AddImage(img)
if not status == openem.kSuccess:
raise RuntimeError("Failed to add image for processing!")
# Process the loaded image.
detections = openem.VectorVectorDetection()
status = detector.Process(detections)
if not status == openem.kSuccess:
raise RuntimeError("Failed to process image {}!".format(img_path))
# Write detection to dict.
for dets in detections:
for det in dets:
path, f = os.path.split(img_path)
frame, _ = os.path.splitext(f)
video_id = os.path.basename(os.path.normpath(path))
x, y, w, h = det.location
if det.confidence >= threshold:
det_data['video_id'].append(video_id)
det_data['frame'].append(frame)
det_data['x'].append(x)
det_data['y'].append(y)
det_data['w'].append(w)
det_data['h'].append(h)
det_data['det_conf'].append(det.confidence)
det_data['det_species'].append(det.species)
print("Finished detection on {}".format(img_path))
# Write detections to csv.
os.makedirs(config.inference_dir(), exist_ok=True)
d = pd.DataFrame(det_data)
d.to_csv(config.detect_inference_path(), index=False)
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)
