SuturePlacer.py

import random

import DistanceCalculator
import RewardFunction
import Constraints
import scipy.optimize as optim
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from datetime import datetime
import os
import math


class SuturePlacer:
    def __init__(self, wound_width, mm_per_pixel, centroids=None):
        # This object should contain the optimizer, the spline curve, the image, etc., i.e. all of the relevant objects involved, as attributes.
        self.wound_width = wound_width
        self.mm_per_pixel = mm_per_pixel
        self.DistanceCalculator = DistanceCalculator.DistanceCalculator(self, self.wound_width, self.mm_per_pixel)
        self.RewardFunction = RewardFunction.RewardFunction(wound_width, self)
        self.Constraints = Constraints.Constraints(wound_width, centroids=centroids)
        self.Constraints.DistanceCalculator = self.DistanceCalculator

        self.b_insert_pts = []
        self.b_center_pts = []
        self.b_extract_pts = []
        self.b_loss = float('inf')

        self.c_lossMin = 0
        self.c_lossIdeal = 1
        self.c_lossVarCenter = 12
        self.c_lossVarInsExt = 6
        self.c_lossClosure = 15
        self.c_lossShear = 5

    def optimize(self, wound_points, optFrame):
        insert_dists, center_dists, extract_dists, insert_pts, center_pts, extract_pts = self.DistanceCalculator.calculate_distances(wound_points)
        self.RewardFunction.insert_dists = insert_dists
        self.RewardFunction.center_dists = center_dists
        self.RewardFunction.extract_dists = extract_dists

        self.Constraints.wound_points = wound_points
        
        def jac(t):
            return optim.approx_fprime(t, final_loss)

        def final_loss(t):
            self.RewardFunction.insert_dists, self.RewardFunction.center_dists, self.RewardFunction.extract_dists, insert_pts, center_pts, extract_pts = self.DistanceCalculator.calculate_distances(t)    
            self.RewardFunction.wound_points = t
            self.RewardFunction.suture_points = list(zip(insert_pts, center_pts, extract_pts))
            return self.RewardFunction.final_loss(c_lossMin=self.c_lossMin, c_lossIdeal = self.c_lossIdeal, c_lossVarCenter = self.c_lossVarCenter, c_lossVarInsExt=self.c_lossVarInsExt, c_lossClosure = self.c_lossClosure, c_lossShear = self.c_lossShear)

        # continuous progress bar
        self.progress += self.progress_incre
        optFrame.update_progress(self.progress)

        self.progress += self.progress_incre
        optFrame.after(100,optFrame.update_progress,self.progress)

        result = optim.minimize(final_loss, wound_points, constraints = self.Constraints.constraints(), options={"maxiter":200}, method = 'SLSQP', tol=1e-2, jac = jac)
        
        # continuous progress bar
        self.progress += self.progress_incre
        optFrame.update_progress(self.progress)
        
        insert_dists, center_dists, extract_dists, insert_pts, center_pts, extract_pts = self.DistanceCalculator.calculate_distances(result.x)
        
        # continuous progress bar
        self.progress += self.progress_incre
        optFrame.update_progress(self.progress)

        self.insert_pts = insert_pts
        self.center_pts = center_pts
        self.extract_pts = extract_pts

        self.progress += self.progress_incre
        optFrame.after(100,optFrame.update_progress,self.progress)

        result = optim.minimize(final_loss, wound_points, constraints = self.Constraints.constraints(), options={"maxiter":200}, method = 'SLSQP', tol = 1e-2, jac = jac)
        plt.clf()
        #save_intermittent_plots = True
        save_intermittent_plots = False
        if save_intermittent_plots:
            self.DistanceCalculator.plot(result.x, "closure", plot_type='closure', save_fig='s1/' + str(len(wound_points)) + '_closure_' + str(random.randint(0, 1000000)))
            self.DistanceCalculator.plot(result.x, "shear", plot_type='shear', save_fig='s1/' + str(len(wound_points)) + '_shear_' + str(random.randint(0, 1000000)))

        return insert_dists, center_dists, extract_dists, insert_pts, center_pts, extract_pts, result.x

    def place_sutures(self, _optFrame, save_figs=False): #save_figs=True

        # make a folder to store info

        # if save_figs:
        #     if not os.path.isdir("clicking"):
        #         os.mkdir('clicking')
            
        #     now = datetime.now()
        #     # dd/mm/YY H:M:S
        #     dt_string = now.strftime("%d-%m-%Y-%H-%M-%S")
        #     os.mkdir('clicking/' + dt_string)
        #     os.mkdir('clicking/' + dt_string + '/sutures')
        #     os.mkdir('clicking/' + dt_string + '/closure')
        #     os.mkdir('clicking/' + dt_string + '/shear')

        num_sutures_initial = int(self.DistanceCalculator.initial_number_of_sutures(0, 1)) # heuristic
        num_sutures_initial = int(num_sutures_initial / 4) # changed suture width calculations from drawing
        print("NUM SUTURES INITIAL:", num_sutures_initial)
        
        # Set up suture range for progress tracking
        start_range = max(2, int(num_sutures_initial))
        end_range = int(2.2 * num_sutures_initial)

        # set up optimization frame
        _optFrame.set_suture_range(start_range,end_range)
        _optFrame.set_distance_calculator(self.DistanceCalculator)
        
        d = {}
        losses = {}
        points_dict = {}
        self.progress_incre = (1 / (end_range - start_range + 1)) / 10
        _optFrame.start_range = start_range
        _optFrame.end_range = end_range
        _optFrame.final_suture_colors = {}
        for num_sutures in range(start_range, end_range): # This should be (0.8 * heuristic to 1.4 * heuristic)
            print('TESTING NUM SUTURES: ', num_sutures)
            
            # Update progress GUI
            _optFrame.update_cur_sutures(num_sutures)
            self.progress = (num_sutures - start_range) / (end_range - start_range + 1)
            
            d[num_sutures] = {}
            heuristic = num_sutures
            best_loss = float('inf')
            wound_points = np.linspace(0, 1, num_sutures)
            
            insert_dists, center_dists, extract_dists, insert_pts, center_pts, extract_pts, ts = self.optimize(wound_points=wound_points,optFrame=_optFrame)
            
            final_suture_colors = ['black'] * num_sutures
            high_curv_sutures = [0] * num_sutures
            closest_sutures = self.Constraints.get_closest_suture(ts)
            for i in closest_sutures:
                final_suture_colors[i] = '#01fd00'
                high_curv_sutures[i] = 1
            _optFrame.final_suture_colors[num_sutures] = final_suture_colors
            _optFrame.high_curv_sutures = high_curv_sutures
            _optFrame.num_high_curv_sutures = np.sum(high_curv_sutures)

            # save all suture plans for later mapping
            _optFrame.planned_insert_pts.append(insert_pts)
            _optFrame.planned_center_pts.append(center_pts)
            _optFrame.planned_extract_pts.append(extract_pts)

            # adaptive length sutures
            dists_to_wound = []
            for cpt in center_pts:
                min_dist = 10000
                scale_len = 0
                #find point in centerline that is closest to center point, save distance to outside of wound
                for opt in _optFrame.parent_root.ordered_pts_dist:
                    temp_dist = math.sqrt((opt[0]-cpt[1])**2 + (opt[1]-cpt[0])**2)
                    if temp_dist < min_dist:
                        min_dist = temp_dist
                        scale_len = opt[2]
                dists_to_wound.append(scale_len)

            def extend_sutures(insert_pts, extract_pts, center_pts, dists_to_wound):
                n_insert_pts = []
                n_extract_pts = []
                for i in range(len(center_pts)):
                    vect = (extract_pts[i][0]-insert_pts[i][0], extract_pts[i][1]-insert_pts[i][1])
                    scale_factor = dists_to_wound[i]*0.1
                    if scale_factor < 1.0:
                        scale_factor = 1.0
                    new_pt = (insert_pts[i][0] + ((scale_factor) * vect[0]), insert_pts[i][1] + ((scale_factor) * vect[1]))
                    n_insert_pts.append(new_pt)
                    new_pt = (extract_pts[i][0] - ((scale_factor) * vect[0]), extract_pts[i][1] - ((scale_factor) * vect[1]))
                    n_extract_pts.append(new_pt)
                
                return n_extract_pts, n_insert_pts
            
            n_insert_pts, n_extract_pts = extend_sutures(insert_pts, extract_pts, center_pts, dists_to_wound)
            _optFrame.planned_n_insert_pts.append(n_insert_pts)
            _optFrame.planned_n_extract_pts.append(n_extract_pts)

            # continuous progress bar
            self.progress += self.progress_incre
            _optFrame.update_progress(self.progress)

            self.RewardFunction.insert_dists = insert_dists
            self.RewardFunction.center_dists = center_dists
            self.RewardFunction.extract_dists = extract_dists
            best_loss = self.RewardFunction.hyperLoss()

            # continuous progress bar
            self.progress += self.progress_incre
            _optFrame.update_progress(self.progress)
            
            # Get individual loss components
            closure_loss = self.RewardFunction.lossClosureForce(1, 0)
            shear_loss = self.RewardFunction.lossClosureForce(0, 1)

            # continuous progress bar
            self.progress += self.progress_incre
            _optFrame.update_progress(self.progress)

            center_var_loss = self.RewardFunction.lossVar(1, 0)
            ins_ext_var_loss = self.RewardFunction.lossVar(0, 1)
            ideal_loss = self.RewardFunction.lossIdeal()

            # continuous progress bar
            self.progress += self.progress_incre
            _optFrame.update_progress(self.progress)
            
            print('loss: ', best_loss)
            print('closure loss', closure_loss)
            print('shear loss', shear_loss)
            # print('center var loss', center_var_loss)
            # print('InsExt var loss', ins_ext_var_loss)
            # print('ideal loss', ideal_loss)

            # save all losses for later plotting
            _optFrame.total_array.append(best_loss)
            _optFrame.closure_array.append(closure_loss)
            _optFrame.shear_array.append(shear_loss)
            
            # Update progress GUI with loss information
            _optFrame.update_losses(best_loss,closure_loss,shear_loss)
            # _optFrame.update_visualization(ts, f'Suture Plan: Optimized {num_sutures} Sutures')
            _optFrame.update_visualization(ts, f'Suture Plan: {num_sutures} Sutures')

            d[num_sutures]['loss'] = best_loss
            d[num_sutures]['closure loss'] = closure_loss
            d[num_sutures]['shear loss'] = shear_loss
            d[num_sutures]['var loss - center'] = center_var_loss
            d[num_sutures]['var loss - ins/ext'] = ins_ext_var_loss
            d[num_sutures]['ideal loss'] = ideal_loss
            b_insert_pts, b_center_pts, b_extract_pts, b_ts = insert_pts, center_pts, extract_pts, ts
            losses[best_loss] = num_sutures
     
            self.insert_pts = b_insert_pts
            self.center_pts = b_center_pts
            self.extract_pts = b_extract_pts

            if best_loss < self.b_loss:
                self.b_loss = best_loss
                self.b_insert_pts = b_insert_pts
                self.b_center_pts = b_center_pts
                self.b_extract_pts = b_extract_pts

            # print(losses)

            # if save_figs:
            
            #     self.DistanceCalculator.plot(b_ts, "Number of Sutures: " + str(num_sutures) + ". Total loss: " + str(best_loss), save_fig=str(num_sutures), plot_type='sutures',save_dir='clicking/'+dt_string)
            #     self.DistanceCalculator.plot(b_ts, "Closure force for " + str(num_sutures) + " sutures", save_fig= str(num_sutures), plot_type='closure', save_dir='clicking/'+dt_string)
            #     self.DistanceCalculator.plot(b_ts, "Shear force for " + str(num_sutures) + " sutures", save_fig=str(num_sutures), plot_type='shear', save_dir='clicking/'+dt_string)

            points_dict[num_sutures] = b_ts

        # Mark optimization as complete
        _optFrame.mark_complete()
        
        dict_to_csv(d, "clicked_losses")
        save_dict_to_file(points_dict, "clicked_points.txt")
        return b_insert_pts, b_center_pts, b_extract_pts

def save_dict_to_file(dic, filename):
    f = open(filename,'w')
    f.write(str(dic))
    f.close()

def load_dict_from_file():
    f = open('dict.txt','r')
    data=f.read()
    f.close()
    return eval(data)

def dict_to_csv(d, filename):
    rows = []
    for k, v in d.items():
        row = {"num_sutures": k}
        row.update(v)
        rows.append(row)
    
    df = pd.DataFrame(rows, columns=['num_sutures', 'loss', 'closure loss', 'shear loss', 'var loss'])
    df = df.sort_values(by=['loss'])
    df.to_csv(filename + ".csv", index=False)