import dlib
import numpy as np
import menpo.io as mio
from menpo.shape import PointCloud, PointDirectedGraph
from menpo.landmark import ibug_face_68

def pointgraph_to_parts(pcloud):
    return [dlib.point(int(p[1]), int(p[0])) for p in pcloud.points]

def full_object_detection_to_pointcloud(full_object_detection):
    return PointCloud(np.vstack([(p.y, p.x) for p in full_object_detection.parts()]))

def detection_to_pointgraph(detection):
    return PointCloud(np.array([(p.y, p.x) for p in detection.parts()]))

def rect_to_pointgraph(rect):
    return bounding_box((rect.top(), rect.left()),
                        (rect.bottom(), rect.right()))

def pointgraph_to_rect(pg):
    min_p, max_p = pg.bounds()
    return dlib.rectangle(left=int(min_p[1]), top=int(min_p[0]),
                          right=int(max_p[1]), bottom=int(max_p[0]))

def bounding_box(min_point, max_point):
    return PointDirectedGraph(
        np.array([min_point, [max_point[0], min_point[1]],
                  max_point, [min_point[0], max_point[1]]]),
        np.array([[0, 1], [1, 2], [2, 3], [3, 0]]), copy=False)

images = list(mio.import_images('lfpw/trainset', max_images=10))

image_pixels = [np.array(im.as_PILImage()) for im in images]

rectangles = [[pointgraph_to_rect(lgroup.lms.bounding_box())
               for lgroup in im.landmarks.values()]
               for im in images]

detections = [[dlib.full_object_detection(pointgraph_to_rect(lgroup.lms.bounding_box()), pointgraph_to_parts(lgroup.lms))
               for lgroup in im.landmarks.values()]
               for im in images]

detection_images_xml = '../examples/faces/training.xml'
prediction_images_xml = '../examples/faces/training_with_face_landmarks.xml'

options = dlib.simple_object_detector_training_options()
options.be_verbose = True

dlib.train_simple_object_detector(detection_images_xml, 'file_detector_test.svm', options)

file_detector = dlib.simple_object_detector('file_detector_test.svm')
print(file_detector)

%matplotlib inline

rects = file_detector(image_pixels[0])

pg = rect_to_pointgraph(rects[0])

im = images[0].copy()
im.landmarks.clear()
im.landmarks['bb'] = pg

im.view_widget()

options = dlib.simple_object_detector_training_options()
options.be_verbose = True

in_memory_detector = dlib.train_simple_object_detector(image_pixels, rectangles, options)

print(in_memory_detector)

%matplotlib inline

rects = in_memory_detector(image_pixels[0])

pg = rect_to_pointgraph(rects[0])
im.landmarks.clear()
im.landmarks['bb'] = pg

im.view_widget()

bundled_detector = dlib.get_frontal_face_detector()

%matplotlib inline

rects = bundled_detector(image_pixels[0])

pg = rect_to_pointgraph(rects[0])
im.landmarks.clear()
im.landmarks['bb'] = pg

im.view_widget()

bundled_detector.save('bundled_detector.svm')
in_memory_detector.save('in_memory_detector.svm')
file_detector.save('file_detector.svm')

bundled_detector = dlib.fhog_object_detector('bundled_detector.svm')
print(bundled_detector)
in_memory_detector = dlib.simple_object_detector('in_memory_detector.svm')
print(in_memory_detector)
file_detector = dlib.simple_object_detector('file_detector.svm')
print(file_detector)

# simple_object_detector - cached upsampling amount
print(dlib.test_simple_object_detector(detection_images_xml, 'file_detector.svm'))
# simple_object_detector - explicit upsampling amount
print(dlib.test_simple_object_detector(detection_images_xml, 'file_detector.svm', upsampling_amount=0))

# fhog_object_detector - default 0 upsampling amount
print(dlib.test_simple_object_detector(detection_images_xml, 'bundled_detector.svm'))
# fhog_object_detector - explicit upsampling amount
print(dlib.test_simple_object_detector(detection_images_xml, 'bundled_detector.svm', upsampling_amount=1))

# simple_object_detector - cached upsampling amount
print(dlib.test_simple_object_detector(image_pixels, rectangles, in_memory_detector))
# simple_object_detector - explicit upsampling amount
print(dlib.test_simple_object_detector(image_pixels, rectangles, in_memory_detector, upsampling_amount=1))

# fhog_object_detector - default 0 upsampling amount
print(dlib.test_simple_object_detector(image_pixels, rectangles, bundled_detector))
# fhog_object_detector - explicit upsampling amount
print(dlib.test_simple_object_detector(image_pixels, rectangles, bundled_detector, upsampling_amount=1))

win = dlib.image_window(image_pixels[0])
win.add_overlay(rects)

# View in green
win2 = dlib.image_window(image_pixels[0])
win2.add_overlay(rects, color=dlib.rgb_pixel(0, 255, 0))

options = dlib.shape_predictor_training_options()
options.be_verbose = True

dlib.train_shape_predictor(prediction_images_xml, 'file_predictor.dat', options)

file_predictor = dlib.shape_predictor('file_predictor.dat')
print(file_detector)

%matplotlib inline

prediction = file_predictor(image_pixels[0], rectangles[0][0])

pc = full_object_detection_to_pointcloud(prediction)
im.landmarks.clear()
im.landmarks['prediction'] = pc

im.view_widget()

options = dlib.shape_predictor_training_options()
options.be_verbose = True

in_memory_predictor = dlib.train_shape_predictor(image_pixels, detections, options)

print(in_memory_predictor)

%matplotlib inline

prediction = in_memory_predictor(image_pixels[0], rectangles[0][0])

pc = full_object_detection_to_pointcloud(prediction)
im.landmarks.clear()
im.landmarks['prediction'] = pc

im.view_widget()

in_memory_predictor.save('in_memory_predictor.dat')
file_predictor.save('file_predictor.dat')

in_memory_predictor = dlib.shape_predictor('in_memory_predictor.dat')
print(in_memory_predictor)
file_predictor = dlib.shape_predictor('file_predictor.dat')
print(file_predictor)

# Test shape predictor
print(dlib.test_shape_predictor(prediction_images_xml, 'file_predictor.dat'))

def interocular_distance(lmark_group):
    lmarks = ibug_face_68(lmark_group)[1]
    l_eye_center = np.mean(lmarks['left_eye'].points)
    r_eye_center = np.mean(lmarks['right_eye'].points)
    return np.linalg.norm(l_eye_center - r_eye_center)

def get_interocular_distances(images):
    distances = []
    for im in images:
        per_image = []
        for group in im.landmarks:
            per_image.append(interocular_distance(im.landmarks[group]))
        distances.append(per_image)
    return distances

# Test shape predictor
print(dlib.test_shape_predictor(image_pixels, detections, in_memory_predictor))

# Test shape predictor with scales
print(dlib.test_shape_predictor(image_pixels, detections, get_interocular_distances(images), in_memory_predictor))

win = dlib.image_window(image_pixels[0])
win.add_overlay(prediction)

# View in green
win2 = dlib.image_window(image_pixels[0])
win2.add_overlay(prediction, color=dlib.rgb_pixel(0, 255, 0))