from numpy import median
from skimage.morphology import square
from matplotlib import pyplot as plt
import cv2

corePath = "/home/piotr/PycharmProjects/Images/samolot"
fileExtension = ".jpg"
images = list()
colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0), (0, 255, 255), (255, 0, 255), (255, 255, 0), (0, 128, 255), (255, 0, 128)]

numberOfColors = len(colors)
smallAreaLimit = 3000
radius = 10

def determineFileNames():
    for i in range(0, 10):
        if i == 6: continue
        images.append(corePath + '0' + str(i) + fileExtension)

    for i in range(10, 19):
        images.append(corePath + str(i) + fileExtension)

def determineContours(fileName):

    image = cv2.imread(fileName, cv2.IMREAD_GRAYSCALE)

    med = median(image)
    image = cv2.erode(image, square(2), iterations=1)

    sigma = 0.6
    lower = int(max(0, (1.0 - sigma) * med))
    upper = int(min(255, (1.0 + sigma) * med))
    image = cv2.Canny(image, lower, upper)

    image = cv2.dilate(image, square(3), iterations=1)
    image = cv2.dilate(image, square(2), iterations=2)

    _, contours, hierarchy = cv2.findContours(image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    return contours, hierarchy

def drawAirplaneContours(fileName, contours, hierarchy):
    image = cv2.imread(fileName, cv2.IMREAD_COLOR)

    for i in range(0, len(contours)):
        if cv2.contourArea(contours[i]) > smallAreaLimit and hierarchy[0][i][3] == -1:
            moments = cv2.moments(contours[i])
            cx, cy = int(moments['m10'] / moments['m00']), int(moments['m01'] / moments['m00'])
            cv2.circle(image, (cx, cy), radius, (255, 255, 255), thickness=3)

            cv2.drawContours(image, [contours[i]], 0, colors[i % numberOfColors], 2)

    return image

def main():
    determineFileNames()

    for picturePath in images:
        contours, hierarchy = determineContours(picturePath)
        image = drawAirplaneContours(picturePath, contours, hierarchy)

        cv2.imwrite(corePath + "Contours" + picturePath[-6:], image)

if __name__ == '__main__':
    main()