Skip to content
Snippets Groups Projects
Commit 3bacfd6d authored by Herman Andersen Dyrkorn's avatar Herman Andersen Dyrkorn
Browse files

Merge branch '16-rework-the-salamander-straightening' into 'master' 

Resolve "Rework the salamander straightening."

Closes #16

See merge request !19
parents 8bdc5d65 cd145269
No related branches found
No related tags found
1 merge request!19Resolve "Rework the salamander straightening."
Show whitespace changes
Inline Side-by-side
algorithm/train_src/new_straightening.py
+ 224
− 134
View file @ 3bacfd6d
......@@ -3,23 +3,13 @@ import numpy as np
from scipy.interpolate import interp1d
import time
import math
import predict_salamander_abdomen as psa
from threading import Semaphore
DEEPLABCUT_CONFIG_PATH= "dlc_model/config.yaml"
# Disse var bruke for testing:
# import os
# directory_path="../img_analyze"
# image_to_predict = "img83.png"
# config = os.path.abspath(DEEPLABCUT_CONFIG_PATH)
# os_directory_path = os.path.abspath(directory_path)
# image_type = image_to_predict[-4:]
# img_directory = os_directory_path + '\\' + image_to_predict
# smallpath = os_directory_path + '/' + 'small' + image_to_predict
# print(smallpath)
from tensorflow.python.keras.layers import Minimum
import predict_salamander_abdomen as psa
from threading import Semaphore
DEEPLABCUT_CONFIG_PATH = "./dlc_model/config.yaml"
########################################
# CONSTANTS
########################################
......@@ -27,7 +17,7 @@ DEEPLABCUT_CONFIG_PATH= "dlc_model/config.yaml"
# What GPU to use (normally 0 if you only have 1 gpu):
GPUID = 0
MAX_DEEPLABCUT_IMAGE_SIZE = 1280
LEAST_DEEPLABCUT_SCORE = 0.01
LEAST_DEEPLABCUT_SCORE = 0.02
# Only one thread is allowed to use tf and the gpu at once:
_ACCESS_TF_AND_GPU_SEMA = Semaphore(1)
......@@ -36,94 +26,170 @@ _ACCESS_TF_AND_GPU_SEMA = Semaphore(1)
STRAIGTHENED_IMAGE_WIDTH = 320
STRAIGTHENED_IMAGE_HEIGHT = 120
STRAIGTHENED_IMAGE_ASPECT_RATIO = STRAIGTHENED_IMAGE_HEIGHT / STRAIGTHENED_IMAGE_WIDTH
# amount of pixels to add to the width of the shoulder including what the AI estimates:
SHOULDER_ESTIMATION_BUFFER = 15
# minimum distance between shoulder points in pixels:
MINIMUM_SHOULDER_WIDTH = 100
# minimum distance between the points, 2 and 3 in pixels:
MINIMUM_MID_POINT_DISTANCE = 30
########################################
# FUNCTIONS
########################################
def straighten(image):
"""
Takes the image and tries to find the abdomen of a salamander by locating
4 points down the abdomen of the salamander. Then we make a curve that
follows those 4 points. Each point will be given a score that tells how
confident the AI is with its estimation.
'''
Takes the image and tries to find the abdomen by locating
4 points down the abdomen of the salamander. Each point will
be given a score that tells how confident the AI with its estimation.
Parameters:
----------
image: cv2 image NEEDS to be RGB not BGR
Full path of the config.yaml file as a string.
----------
Returns:
-------
straightened_image: The straightened image of the abdomen. None if straightening fails.
If there is one score that falls below a given constant the function
return None, Scores
cropped_image: A potentially cropped version of the original image.
Returns: image. A straightened image.
score. A list (length = 4) of scores for each estimation
estimated
'''
def straighten(image):
points_spine: The spine points that were found by the AI.
points_shoulder: The shoulder points that were found by the AI.
score: How confident the AI was at each point.
-------
"""
with _ACCESS_TF_AND_GPU_SEMA:
start = time.time()
start_time = time.time()
height, width, _ = image.shape
# find the largest dimention:
larger_dim = max(width, height)
img = image
cropped_image = image
# reduce size of image by the largest dimention, if it exceeds a given constant:
if larger_dim > MAX_DEEPLABCUT_IMAGE_SIZE:
factor = MAX_DEEPLABCUT_IMAGE_SIZE / larger_dim
img = cv2.resize(img, None, fx=factor, fy=factor, interpolation=cv2.INTER_CUBIC)
cropped_image = cv2.resize(cropped_image, None, fx=factor, fy=factor, interpolation=cv2.INTER_CUBIC)
prediction, _, x, y = psa.get_prediction_dlc(config=DEEPLABCUT_CONFIG_PATH, image=cropped_image, gputouse=GPUID)
prediction, _, _, _ = psa.get_prediction_dlc(config=DEEPLABCUT_CONFIG_PATH, image=img, gputouse=GPUID)
# points is an array of coordinates to use to straighten the image:
# We use the prediction from deeplabcut for the points. There are two things to take note of.
# The predictions can be negative, so we ensure that negative values are set as 0 with max(0,num)
# the structure of prediction is: shape(1,12) content: [bp1_x, bp1_y, bp1_score, bp2_x, bp2_y, bp2_score
# bp3_x, bp3_y, bp3_score, bp4_x, bp4_y, bp4_score] 'bp' = bodypart.
# we are using the score to determine if the image is valid and skipping the score for points:
score = np.array([prediction[2],prediction[5],prediction[8],prediction[11]])
score = np.array([prediction[2], prediction[5], prediction[8], prediction[11], prediction[14], prediction[17]])
if score[0] < LEAST_DEEPLABCUT_SCORE or score[1] < LEAST_DEEPLABCUT_SCORE or \
score[2] < LEAST_DEEPLABCUT_SCORE or score[3] < LEAST_DEEPLABCUT_SCORE:
print("straighten failed because the score (", score, " found is below ", LEAST_DEEPLABCUT_SCORE)
return None, None, None, score
return None, None, None, None, score
# getting new width and height:
heigh, width, _ = img.shape
heigh, width, _ = cropped_image.shape
# ensureing each point is between 0 and width/height:
points = np.array([
[min(width,max(0,int(prediction[0]))),min(heigh,max(0,int(prediction[1])))],
[min(width, max(0,int(prediction[3]))),min(heigh,max(0,int(prediction[4])))],
[min(width,max(0,int(prediction[6]))),min(heigh, max(0,int(prediction[7])))],
[min(width,max(0,int(prediction[9]))),min(heigh, max(0,int(prediction[10])))],
points_spine = np.array([
(min(width, max(0, int(prediction[0]))), min(heigh, max(0, int(prediction[1])))),
(min(width, max(0, int(prediction[3]))), min(heigh, max(0, int(prediction[4])))),
(min(width, max(0, int(prediction[6]))), min(heigh, max(0, int(prediction[7])))),
(min(width, max(0, int(prediction[9]))), min(heigh, max(0, int(prediction[10]))))
])
# Used to check if the point 2 is closer to point 1. This will also be used later
# on to determine if the two points in the middle have been swapped.
dist_2_1 = np.linalg.norm(points_spine[1] - points_spine[0])
# print("dist_2_1: ", dist_2_1)
dist_2_3 = np.linalg.norm(points_spine[1] - points_spine[2])
# check if the two middle points are equal:
if dist_2_3 < MINIMUM_MID_POINT_DISTANCE:
print("the middle points needs to be moved")
dist_2_4 = np.linalg.norm(points_spine[1] - points_spine[3])
# print("dist_2_4: ", dist_2_4)
# if point 2 needs to move towards the 1 point:
if dist_2_1 > dist_2_4:
print("moving the 2 point torwards the 1")
points_spine[1] = halfway_between(point1=points_spine[1], point2=points_spine[0])
# else the 3 point needs to be moved towards the 4 point:
else:
print("moving the 3 point torwards the 4")
points_spine[2] = halfway_between(point1=points_spine[2], point2=points_spine[3])
del dist_2_3
# Shoulder points are cannot be a part of the spine points as the spide points are sednt
# to a function that expects all the points to represent a line:
points_shoulder = np.array([
(min(width, max(0, int(prediction[12]))), min(heigh, max(0, int(prediction[13])))),
(min(width, max(0, int(prediction[15]))), min(heigh, max(0, int(prediction[16]))))
])
######### VALIDATING SPINE POINTS:
# This would mean that either the first or last point are above each other.
# or the image is really small. This is ugly, but we can't just check for duplicates:
if np.linalg.norm(points_spine[0] - points_spine[1]) < MINIMUM_MID_POINT_DISTANCE or \
np.linalg.norm(points_spine[0] - points_spine[2]) < MINIMUM_MID_POINT_DISTANCE or \
np.linalg.norm(points_spine[0] - points_spine[3]) < MINIMUM_MID_POINT_DISTANCE or \
np.linalg.norm(points_spine[3] - points_spine[1]) < MINIMUM_MID_POINT_DISTANCE or \
np.linalg.norm(points_spine[3] - points_spine[2]) < MINIMUM_MID_POINT_DISTANCE:
print("predicted points were not correct"
" and we can't proceed with the straightening")
return None, None, None, None, score
# checking if the points in the middle needs to be swapped:
dist_from_2_to_1 = math.sqrt((points[1][0] - points[0][0]) ** 2 + (points[1][1] - points[0][1]) ** 2)
dist_from_3_to_1 = math.sqrt((points[2][0] - points[0][0]) ** 2 + (points[2][1] - points[0][1]) ** 2)
print(points)
if dist_from_2_to_1 > dist_from_3_to_1:
points[[1, 2], :] = points[[2, 1], :]
# points[1], points[2] = points[2], points[1]
print("heihei", dist_from_3_to_1, " ", dist_from_2_to_1)
del dist_from_3_to_1
del dist_from_2_to_1
print(points)
curve = get_smooth_curve(points, STRAIGTHENED_IMAGE_WIDTH)
dist_2_1 = np.linalg.norm(points_spine[1] - points_spine[0])
dist_3_1 = np.linalg.norm(points_spine[2] - points_spine[0])
if dist_2_1 > dist_3_1:
points_spine[[1, 2], :] = points_spine[[2, 1], :]
print("Two middle points need to be swapped, ", dist_3_1, " ", dist_2_1)
del dist_3_1
del dist_2_1
# Finding shoulder width because the pattern on the salamander is usually as wide as
# its shoulders. This will be used to stretch the pattern to the image borders. If
# we don't do this a lot of the image will be of either the salamanders side (that
# does not have a pattern) or background:
shoulder_width = math.sqrt((points_shoulder[1][0] - points_shoulder[0][0]) ** 2 \
+ (points_shoulder[1][1] - points_shoulder[0][1]) ** 2)
print("shoulder width: ", shoulder_width)
if shoulder_width < MINIMUM_SHOULDER_WIDTH:
shoulder_width = 2 * (math.sqrt((points_shoulder[1][0] - points_spine[0][0]) ** 2 \
+ (points_shoulder[1][1] - points_spine[0][1]) ** 2))
print("correcting: ", shoulder_width)
if shoulder_width < MINIMUM_SHOULDER_WIDTH:
shoulder_width = STRAIGTHENED_IMAGE_HEIGHT
start = time.time()
curve = get_smooth_curve(points_spine, STRAIGTHENED_IMAGE_WIDTH)
end = time.time()
print("Bicubic interpolation of spine took " + str(end - start) + "s")
start = time.time()
map = generate_map_from_bellycurve(curve)
result = cv2.remap(img, map[:, :, 0], map[:, :, 1], cv2.INTER_LINEAR)
map = generate_map_from_bellycurve(curve, shoulder_width)
straightened_image = cv2.remap(cropped_image, map[:, :, 0], map[:, :, 1], cv2.INTER_LINEAR)
end_time = time.time()
print("Straightening took " + str(end_time - start_time) + "s")
return straightened_image, cropped_image, points_spine, points_shoulder, score
end = time.time()
print("Straightening took " + str(end - start) + "s")
return result, img, points, score
def generate_map_from_bellycurve(curve, width=STRAIGTHENED_IMAGE_WIDTH):
'''
Generates a map that can be passed to cv2.remap to extract the belly pattern
'''
def generate_map_from_bellycurve(curve):
salamander_length = np.linalg.norm(curve[len(curve) - 1] - curve[0])
salamander_width = salamander_length * STRAIGTHENED_IMAGE_ASPECT_RATIO
# salamander_width = salamander_length * STRAIGTHENED_IMAGE_ASPECT_RATIO
salamander_width = width + SHOULDER_ESTIMATION_BUFFER
gradient = np.gradient(curve, axis=0)
gradientlength = np.linalg.norm(gradient, axis=1)
# print("curve: ", curve.shape)
# print("gradient: ", gradient.shape)
# print("gradientlength: ", gradientlength.shape)
tu = np.divide(gradient, gradientlength.reshape((STRAIGTHENED_IMAGE_WIDTH, 1)))
su = np.ndarray(tu.shape)
......@@ -139,12 +205,16 @@ def generate_map_from_bellycurve(curve):
return map.astype('float32')
def get_smooth_curve(points, num_points):
'''
Takes in an array of 2-D points and returns a new set of points
with numpoints elements where the missing points are interpolated
using cubic interpolation
'''
def get_smooth_curve(points, num_points):
# Calculate distance between all the points
distance = np.sqrt(np.sum(np.diff(points, axis=0) ** 2, axis=1))
# print("distance, ", distance.shape)
......@@ -168,3 +238,23 @@ def get_smooth_curve(points, num_points):
interpoints = f(alpha)
# print("interpoint, ", interpoints.shape)
return interpoints
# def get_distance(point1, point2):
# '''
# get_distance() calculates the distance between two points:
# '''
# return 2 * (math.sqrt((point1[0] - point2[0]) ** 2 \
# + (point1[1] - point2[1]) ** 2))
#
def halfway_between(point1, point2):
"""
halfway_between() calculates a point between two points:
"""
vec = [(point2[0] - point1[0]) / 2, (point2[1] - point1[1]) / 2]
return [point1[0] + vec[0], point1[1] + vec[1]]
algorithm/train_src/test_straightening.py
+ 14
− 14
View file @ 3bacfd6d
......@@ -21,13 +21,10 @@ img = cv2.imread(img_directory)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
str_image, img, points, _ = straighten.straighten(image = img)
str_image = cv2.cvtColor(str_image, cv2.COLOR_BGR2RGB)
# data points
str_image, img, points, shoulder_points, _ = straighten.straighten(image = img)
if str_image is not None:
str_image = cv2.cvtColor(str_image, cv2.COLOR_BGR2RGB)
cv2.imwrite(directory_path + '/' + image_to_predict[0:-4] + '_str.png', str_image)
# data points
implot = plt.imshow(img)
print(points[0])
......@@ -36,4 +33,7 @@ plt.scatter(points[0][0], points[0][1], c='b', s=40, )
plt.scatter(points[1][0], points[1][1], c='r', s=40, )
plt.scatter(points[2][0], points[2][1], c='g', s=40, )
plt.scatter(points[3][0], points[3][1], c='y', s=40, )
plt.scatter(shoulder_points[0][0], shoulder_points[0][1], c='y', s=40, )
plt.scatter(shoulder_points[1][0], shoulder_points[1][1], c='y', s=40, )
plt.savefig(directory_path + '/' + image_to_predict[0:-4] + '_points.png')
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment