def check_image(fname, poller):
"""
Check method that ensures the image is valid.
:param fname: the file to check
:type fname: str
:param poller: the Poller instance that called the method
:type poller: Poller
:return: True if complete
:rtype: bool
"""
result = auto.is_image_complete(fname)
poller.debug("Image complete:", fname, "->", result)
return result
def predict_on_images(input_dir, graph, sess, output_dir, tmp_dir, score_threshold, categories, num_imgs, inference_times,
delete_input, output_polygons, mask_threshold, mask_nth, output_minrect, view_margin, fully_connected,
fit_bbox_to_polygon, output_width_height, bbox_as_fallback):
"""
Method performing predictions on all images ony by one or combined as specified by the int value of num_imgs.
:param input_dir: the directory with the images
:type input_dir: str
:param graph: the graph to use
:type graph: tf.Graph()
:param sess: the tensorflow session
:type sess: tf.Session
:param output_dir: the output directory to move the images to and store the predictions
:type output_dir: str
:param tmp_dir: the temporary directory to store the predictions until finished
:type tmp_dir: str
:param score_threshold: the minimum score predictions have to have
:type score_threshold: float
:param categories: the label map
:param num_imgs: the number of images to combine into one before presenting to graph
:type num_imgs: int
:param inference_times: whether to output a CSV file with the inference times
:type inference_times: bool
:param delete_input: whether to delete the input images rather than moving them to the output directory
:type delete_input: bool
:param output_polygons: whether the model predicts masks and polygons should be stored in the CSV files
:type output_polygons: bool
:param mask_threshold: the threshold to use for determining the contour of a mask
:type mask_threshold: float
:param mask_nth: to speed up polygon computation, use only every nth row and column from mask
:type mask_nth: int
:param output_minrect: when predicting polygons, whether to output the minimal rectangles around the objects as well
:type output_minrect: bool
:param view_margin: the margin in pixels to use around the masks
:type view_margin: int
:param fully_connected: whether regions of 'high' or 'low' values should be fully-connected at isthmuses
:type fully_connected: str
:param fit_bbox_to_polygon: whether to fit the bounding box to the polygon
:type fit_bbox_to_polygon: bool
:param output_width_height: whether to output x/y/w/h instead of x0/y0/x1/y1
:type output_width_height: bool
:param bbox_as_fallback: if ratio between polygon-bbox and bbox is smaller than this value, use bbox as fallback polygon, ignored if < 0
:type bbox_as_fallback: float
"""
# Iterate through all files present in "test_images_directory"
total_time = 0
if inference_times:
times = list()
times.append("Image(s)_file_name(s),Total_time(ms),Number_of_images,Time_per_image(ms)\n")
# counter for keeping track of images that cannot be processed
incomplete_counter = dict()
while True:
start_time = datetime.now()
im_list = []
# Loop to pick up images equal to num_imgs or the remaining images if less
for image_path in os.listdir(input_dir):
# Load images only
ext_lower = os.path.splitext(image_path)[1]
if ext_lower in SUPPORTED_EXTS:
full_path = os.path.join(input_dir, image_path)
if auto.is_image_complete(full_path):
im_list.append(full_path)
else:
if not full_path in incomplete_counter:
incomplete_counter[full_path] = 1
else:
incomplete_counter[full_path] = incomplete_counter[full_path] + 1
# remove images that cannot be processed
remove_from_blacklist = []
for k in incomplete_counter:
if incomplete_counter[k] == MAX_INCOMPLETE:
print("%s - %s" % (str(datetime.now()), os.path.basename(k)))
remove_from_blacklist.append(k)
try:
if delete_input:
print(" flagged as incomplete {} times, deleting\n".format(MAX_INCOMPLETE))
os.remove(k)
else:
print(" flagged as incomplete {} times, skipping\n".format(MAX_INCOMPLETE))
os.rename(k, os.path.join(output_dir, os.path.basename(k)))
except:
print(traceback.format_exc())
for k in remove_from_blacklist:
del incomplete_counter[k]
if len(im_list) == num_imgs:
break
if len(im_list) == 0:
time.sleep(1)
break
else:
print("%s - %s" % (str(datetime.now()), ", ".join(os.path.basename(x) for x in im_list)))
try:
# Combining picked up images
i = len(im_list)
combined = []
comb_img = None
if i > 1:
while i != 0:
if comb_img is None:
img2 = Image.open(im_list[i-1])
img1 = Image.open(im_list[i-2])
i -= 1
combined.append(os.path.join(output_dir, "combined.png"))
else:
img2 = comb_img
img1 = Image.open(im_list[i-1])
i -= 1
# Remove alpha channel if present
img1 = remove_alpha_channel(img1)
img2 = remove_alpha_channel(img2)
w1, h1 = img1.size
w2, h2 = img2.size
comb_img = np.zeros((h1+h2, max(w1, w2), 3), np.uint8)
comb_img[:h1, :w1, :3] = img1
comb_img[h1:h1+h2, :w2, :3] = img2
comb_img = Image.fromarray(comb_img)
if comb_img is None:
im_name = im_list[0]
image = Image.open(im_name)
image = remove_alpha_channel(image)
else:
im_name = combined[0]
image = remove_alpha_channel(comb_img)
image_np = image_to_numpyarray(image)
output_dict = inference_for_image(image_np, graph, sess)
# Loading results
boxes = output_dict['detection_boxes']
scores = output_dict['detection_scores']
classes = output_dict['detection_classes']
# Code for splitting rois to multiple csv's, one csv per image before combining
max_height = 0
prev_min = 0
for i in range(len(im_list)):
img = Image.open(im_list[i])
img_height = img.height
min_height = prev_min
max_height += img_height
prev_min = max_height
roi_path = "{}/{}-rois.csv".format(output_dir, os.path.splitext(os.path.basename(im_list[i]))[0])
if tmp_dir is not None:
roi_path_tmp = "{}/{}-rois.tmp".format(tmp_dir, os.path.splitext(os.path.basename(im_list[i]))[0])
else:
roi_path_tmp = "{}/{}-rois.tmp".format(output_dir, os.path.splitext(os.path.basename(im_list[i]))[0])
roiobjs = []
for index in range(output_dict['num_detections']):
score = scores[index]
# Ignore this roi if the score is less than the provided threshold
if score < score_threshold:
continue
y0n, x0n, y1n, x1n = boxes[index]
# Translate roi coordinates into combined image coordinates
x0 = x0n * image.width
y0 = y0n * image.height
x1 = x1n * image.width
y1 = y1n * image.height
if y0 > max_height or y1 > max_height:
continue
elif y0 < min_height or y1 < min_height:
continue
label = classes[index]
label_str = categories[label - 1]['name']
px = None
py = None
pxn = None
pyn = None
bw = None
bh = None
if output_polygons:
px = []
py = []
pxn = []
pyn = []
bw = ""
bh = ""
if 'detection_masks'in output_dict:
poly = mask_to_polygon(output_dict['detection_masks'][index], mask_threshold=mask_threshold,
mask_nth=mask_nth, view=(x0, y0, x1, y1), view_margin=view_margin,
fully_connected=fully_connected)
if len(poly) > 0:
px, py = polygon_to_lists(poly[0], swap_x_y=True, normalize=False)
pxn, pyn = polygon_to_lists(poly[0], swap_x_y=True, normalize=True, img_width=image.width, img_height=image.height)
if output_minrect:
bw, bh = polygon_to_minrect(poly[0])
if bbox_as_fallback >= 0:
if len(px) >= 3:
p_x0n, p_y0n, p_x1n, p_y1n = polygon_to_bbox(lists_to_polygon(pxn, pyn))
p_area = (p_x1n - p_x0n) * (p_y1n - p_y0n)
b_area = (x1n - x0n) * (y1n - y0n)
if (b_area > 0) and (p_area / b_area < bbox_as_fallback):
px = [float(i) for i in [x0, x1, x1, x0]]
py = [float(i) for i in [y0, y0, y1, y1]]
pxn = [float(i) for i in [x0n, x1n, x1n, x0n]]
pyn = [float(i) for i in [y0n, y0n, y1n, y1n]]
else:
px = [float(i) for i in [x0, x1, x1, x0]]
py = [float(i) for i in [y0, y0, y1, y1]]
pxn = [float(i) for i in [x0n, x1n, x1n, x0n]]
pyn = [float(i) for i in [y0n, y0n, y1n, y1n]]
if output_minrect:
bw = x1 - x0 + 1
bh = y1 - y0 + 1
if fit_bbox_to_polygon:
if len(px) >= 3:
x0, y0, x1, y1 = polygon_to_bbox(lists_to_polygon(px, py))
x0n, y0n, x1n, y1n = polygon_to_bbox(lists_to_polygon(pxn, pyn))
roiobj = ROIObject(x0, y0, x1, y1, x0n, y0n, x1n, y1n, label, label_str, score=score,
poly_x=px, poly_y=py, poly_xn=pxn, poly_yn=pyn,
minrect_w=bw, minrect_h=bh)
roiobjs.append(roiobj)
info = ImageInfo(os.path.basename(im_list[i]))
roiext = (info, roiobjs)
options = ["--output", str(tmp_dir if tmp_dir is not None else output_dir), "--no-images"]
if output_width_height:
options.append("--size-mode")
roiwriter = ROIWriter(options)
roiwriter.save([roiext])
if tmp_dir is not None:
os.rename(roi_path_tmp, roi_path)
except:
print("Failed processing images: {}".format(",".join(im_list)))
print(traceback.format_exc())
# Move finished images to output_path or delete it
for i in range(len(im_list)):
if delete_input:
os.remove(im_list[i])
else:
os.rename(im_list[i], os.path.join(output_dir, os.path.basename(im_list[i])))
end_time = datetime.now()
inference_time = end_time - start_time
inference_time = int(inference_time.total_seconds() * 1000)
time_per_image = int(inference_time / len(im_list))
if inference_times:
l = ""
for i in range(len(im_list)):
l += ("{}|".format(os.path.basename(im_list[i])))
l += ",{},{},{}\n".format(inference_time, len(im_list), time_per_image)
times.append(l)
print(" Inference + I/O time: {} ms\n".format(inference_time))
total_time += inference_time
if inference_times:
with open(os.path.join(output_dir, "inference_time.csv"), "w") as time_file:
for l in times:
time_file.write(l)
with open(os.path.join(output_dir, "total_time.txt"), "w") as total_time_file:
total_time_file.write("Total inference and I/O time: {} ms\n".format(total_time))
def test_complete_gif(self):
self.assertTrue(is_image_complete(self.data_file("complete.gif")))
def test_complete_jpg(self):
self.assertTrue(is_image_complete(self.data_file("complete.jpg")))
def test_empty_png(self):
self.assertFalse(is_image_complete(self.data_file("empty.png")))
def test_unknown_type(self):
try:
is_image_complete(self.data_file("file.blah"))
self.fail("should have failed with exception")
except:
pass
def test_empty_gif(self):
self.assertFalse(is_image_complete(self.data_file("empty.gif")))
def test_incomplete_png(self):
self.assertFalse(is_image_complete(self.data_file("incomplete.png")))
def test_incomplete_gif(self):
self.assertFalse(is_image_complete(self.data_file("incomplete.gif")))
def predict_on_images(model, input_dir, output_dir, tmp_dir, delete_input, clash_suffix="-in"):
"""
Performs predictions on images found in input_dir and outputs the prediction PNG files in output_dir.
:param model: the model to use
:param input_dir: the directory with the images
:type input_dir: str
:param output_dir: the output directory to move the images to and store the predictions
:type output_dir: str
:param tmp_dir: the temporary directory to store the predictions until finished
:type tmp_dir: str
:param delete_input: whether to delete the input images rather than moving them to the output directory
:type delete_input: bool
:param clash_suffix: the suffix to use for clashes, ie when the input is already a PNG image
:type clash_suffix: str
"""
# counter for keeping track of images that cannot be processed
incomplete_counter = dict()
num_imgs = 1
while True:
start_time = datetime.now()
im_list = []
# Loop to pick up images equal to num_imgs or the remaining images if less
for image_path in os.listdir(input_dir):
# Load images only
ext_lower = os.path.splitext(image_path)[1]
if ext_lower in SUPPORTED_EXTS:
full_path = os.path.join(input_dir, image_path)
if auto.is_image_complete(full_path):
im_list.append(full_path)
else:
if not full_path in incomplete_counter:
incomplete_counter[full_path] = 1
else:
incomplete_counter[full_path] = incomplete_counter[full_path] + 1
# remove images that cannot be processed
remove_from_blacklist = []
for k in incomplete_counter:
if incomplete_counter[k] == MAX_INCOMPLETE:
print("%s - %s" % (str(datetime.now()), os.path.basename(k)))
remove_from_blacklist.append(k)
try:
if delete_input:
print(" flagged as incomplete {} times, deleting\n".format(MAX_INCOMPLETE))
os.remove(k)
else:
print(" flagged as incomplete {} times, skipping\n".format(MAX_INCOMPLETE))
os.rename(k, os.path.join(output_dir, os.path.basename(k)))
except:
print(traceback.format_exc())
for k in remove_from_blacklist:
del incomplete_counter[k]
if len(im_list) == num_imgs:
break
if len(im_list) == 0:
time.sleep(1)
break
else:
print("%s - %s" % (str(datetime.now()), ", ".join(os.path.basename(x) for x in im_list)))
try:
for i in range(len(im_list)):
parts = os.path.splitext(os.path.basename(im_list[i]))
if tmp_dir is not None:
out_file = os.path.join(tmp_dir, parts[0] + ".png")
else:
out_file = os.path.join(output_dir, parts[0] + ".png")
model.predict_segmentation(inp=im_list[i], out_fname=out_file)
except:
print("Failed processing images: {}".format(",".join(im_list)))
print(traceback.format_exc())
# Move finished images to output_path or delete it
for i in range(len(im_list)):
if delete_input:
os.remove(im_list[i])
else:
# PNG input clashes with output, append suffix
if im_list[i].lower().endswith(".png"):
parts = os.path.splitext(os.path.basename(im_list[i]))
os.rename(im_list[i], os.path.join(output_dir, parts[0] + clash_suffix + parts[1]))
else:
os.rename(im_list[i], os.path.join(output_dir, os.path.basename(im_list[i])))
end_time = datetime.now()
inference_time = end_time - start_time
inference_time = int(inference_time.total_seconds() * 1000)
print(" Inference + I/O time: {} ms\n".format(inference_time))
def predict(model,
input_dir,
output_dir,
tmp_dir=None,
top_k=5,
score_threshold=0.0,
delete_input=False,
output_polygons=False,
mask_threshold=0.1,
mask_nth=1,
output_minrect=False,
view_margin=2,
fully_connected='high',
fit_bbox_to_polygon=False,
output_width_height=False,
bbox_as_fallback=False,
scale=1.0,
debayer="",
continuous=False,
output_mask_image=False):
"""
Loads the model/config and performs predictions.
:param model: the model to use
:type model: Yolact
:param input_dir: the directory to check for images to use for prediction
:type input_dir: str
:param output_dir: the directory to store the results in (predictions and/or images)
:type output_dir: str
:param tmp_dir: the directory to store the results initially, before moving them into the actual output directory
:type tmp_dir: str
:param top_k: the top X predictions (= objects) to parse
:type top_k: int
:param score_threshold: the score threshold to use
:type score_threshold: float
:param delete_input: whether to delete the images from the input directory rather than moving them to the output directory
:type delete_input: bool
:param output_polygons: whether the model predicts masks and polygons should be stored in the CSV files
:type output_polygons: bool
:param mask_threshold: the threshold to use for determining the contour of a mask
:type mask_threshold: float
:param mask_nth: to speed up polygon computation, use only every nth row and column from mask
:type mask_nth: int
:param output_minrect: when predicting polygons, whether to output the minimal rectangles around the objects as well
:type output_minrect: bool
:param view_margin: the margin in pixels to use around the masks
:type view_margin: int
:param fully_connected: whether regions of 'high' or 'low' values should be fully-connected at isthmuses
:type fully_connected: str
:param fit_bbox_to_polygon: whether to fit the bounding box to the polygon
:type fit_bbox_to_polygon: bool
:param output_width_height: whether to output x/y/w/h instead of x0/y0/x1/y1
:type output_width_height: bool
:param bbox_as_fallback: if ratio between polygon-bbox and bbox is smaller than this value, use bbox as fallback polygon, ignored if < 0
:type bbox_as_fallback: float
:param scale: the scale to use for the image (0-1)
:type scale: float
:param debayer: the OpenCV2 debayering type to use, eg COLOR_BAYER_BG2BGR; ignored if empty string or None
:type debayer: str
:param continuous: whether to delete the images from the input directory rather than moving them to the output directory
:type continuous: bool
:param output_mask_image: when generating masks, whether to output a combined mask image as well
:type output_mask_image: bool
"""
# counter for keeping track of images that cannot be processed
incomplete_counter = dict()
num_imgs = 1
# evaluate debayering constant
debayer_int = None
if (debayer is not None) and debayer.startswith("COLOR_BAYER_"):
debayer_int = int(eval("cv2." + debayer))
while True:
start_time = datetime.now()
im_list = []
# Loop to pick up images equal to num_imgs or the remaining images if less
for image_path in os.listdir(input_dir):
# Load images only
ext_lower = os.path.splitext(image_path)[1]
if ext_lower in SUPPORTED_EXTS:
full_path = os.path.join(input_dir, image_path)
if auto.is_image_complete(full_path):
im_list.append(full_path)
else:
if not full_path in incomplete_counter:
incomplete_counter[full_path] = 1
else:
incomplete_counter[
full_path] = incomplete_counter[full_path] + 1
# remove images that cannot be processed
remove_from_blacklist = []
for k in incomplete_counter:
if incomplete_counter[k] == MAX_INCOMPLETE:
print("%s - %s" %
(str(datetime.now()), os.path.basename(k)))
remove_from_blacklist.append(k)
try:
if delete_input:
print(
" flagged as incomplete {} times, deleting\n".
format(MAX_INCOMPLETE))
os.remove(k)
else:
print(
" flagged as incomplete {} times, skipping\n".
format(MAX_INCOMPLETE))
os.rename(
k, os.path.join(output_dir,
os.path.basename(k)))
except:
print(traceback.format_exc())
for k in remove_from_blacklist:
del incomplete_counter[k]
if len(im_list) == num_imgs:
break
if len(im_list) == 0:
if continuous:
time.sleep(1)
continue
else:
break
else:
print("%s - %s" % (str(datetime.now()), ", ".join(
os.path.basename(x) for x in im_list)))
try:
for i in range(len(im_list)):
roi_path = "{}/{}-rois.csv".format(
output_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
img_path = "{}/{}-mask.png".format(
output_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
if tmp_dir is not None:
roi_path_tmp = "{}/{}-rois.tmp".format(
tmp_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
img_path_tmp = "{}/{}-mask.tmp".format(
tmp_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
else:
roi_path_tmp = "{}/{}-rois.tmp".format(
output_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
img_path_tmp = "{}/{}-mask.tmp".format(
output_dir,
os.path.splitext(os.path.basename(im_list[i]))[0])
info, roiobjs, mask_comb = predict_image(
model,
im_list[i],
top_k=top_k,
score_threshold=score_threshold,
output_polygons=output_polygons,
mask_threshold=mask_threshold,
mask_nth=mask_nth,
output_minrect=output_minrect,
view_margin=view_margin,
fully_connected=fully_connected,
fit_bbox_to_polygon=fit_bbox_to_polygon,
bbox_as_fallback=bbox_as_fallback,
scale=scale,
debayer_int=debayer_int,
output_mask_image=output_mask_image)
roiext = (info, roiobjs)
options = [
"--output",
str(tmp_dir if tmp_dir is not None else output_dir),
"--no-images"
]
if output_width_height:
options.append("--size-mode")
roiwriter = ROIWriter(options)
roiwriter.save([roiext])
if tmp_dir is not None:
os.rename(roi_path_tmp, roi_path)
if mask_comb is not None:
im = Image.fromarray(np.uint8(mask_comb), 'P')
im.save(img_path_tmp, "PNG")
os.rename(img_path_tmp, img_path)
except:
print("Failed processing images: {}".format(",".join(im_list)))
print(traceback.format_exc())
# Move finished images to output_path or delete it
for i in range(len(im_list)):
if delete_input:
os.remove(im_list[i])
else:
os.rename(
im_list[i],
os.path.join(output_dir, os.path.basename(im_list[i])))
end_time = datetime.now()
inference_time = end_time - start_time
inference_time = int(inference_time.total_seconds() * 1000)
print(" Inference + I/O time: {} ms\n".format(inference_time))