"""
This module contains four napari widgets declared in
different ways:

- a pure Python function flagged with `autogenerate: true`
    in the plugin manifest. Type annotations are used by
    magicgui to generate widgets for each parameter. Best
    suited for simple processing tasks - usually taking
    in and/or returning a layer.
- a `magic_factory` decorated function. The `magic_factory`
    decorator allows us to customize aspects of the resulting
    GUI, including the widgets associated with each parameter.
    Best used when you have a very simple processing task,
    but want some control over the autogenerated widgets. If you
    find yourself needing to define lots of nested functions to achieve
    your functionality, maybe look at the `Container` widget!
- a `magicgui.widgets.Container` subclass. This provides lots
    of flexibility and customization options while still supporting
    `magicgui` widgets and convenience methods for creating widgets
    from type annotations. If you want to customize your widgets and
    connect callbacks, this is the best widget option for you.
- a `QWidget` subclass. This provides maximal flexibility but requires
    full specification of widget layouts, callbacks, events, etc.

References:
- Widget specification: https://napari.org/stable/plugins/guides.html?#widgets
- magicgui docs: https://pyapp-kit.github.io/magicgui/

Replace code below according to your needs.
"""
from typing import TYPE_CHECKING
from napari.utils.notifications import show_info, show_error
from napari.utils import progress
from magicgui import magic_factory
from qtpy.QtWidgets import QHBoxLayout, QPushButton, QWidget, QLineEdit, QLabel, QFileDialog, QVBoxLayout, QGridLayout, QFormLayout, QGroupBox, QDialog, QSpinBox, QListWidget, QListWidgetItem, QComboBox
from qtpy.QtGui import QPalette, QColor
from magicgui.widgets import Table
from skimage import img_as_ubyte
from skimage.io import imread, imsave
from collections import Counter
import numpy as np
import pandas as pd
from PIL import Image
from urllib.request import urlopen
from manini.utils import get_output_size_image, dice_coefficient
import os
import pathlib
import tempfile
import subprocess
import shutil
import sys
from tqdm import tqdm
from manini.script2 import TableWidget,comboCompanies
from pandas import DataFrame
from zipfile import ZipFile        
from magicgui.tqdm import trange
from napari.layers import Layer

from time import process_time
from napari.types import ImageData, LabelsData
if TYPE_CHECKING:
    import napari

class Image_segmentation(QDialog):
    
    def __init__(self, parent: QWidget):
        print("Image segmentation OPEN")
        super().__init__(parent)
        self.setWindowTitle("Image segmentation")
        self.number = QSpinBox()
        
        self.filename_edit = QLineEdit()  
        
        self.file_name_image = QPushButton("File")
        self.filename_edit_image = QLineEdit()  
        self.file_name_image.clicked.connect(self.open_file_dialog_image)
        
        self.file_name_model = QPushButton("File")
        self.filename_edit_model = QLineEdit()  
        self.file_name_model.clicked.connect(self.open_file_dialog_model)

        self.file_name_class_name = QPushButton("File")
        self.filename_edit_class_name = QLineEdit()  
        self.file_name_class_name.clicked.connect(self.open_file_dialog_class_name)
        
        self.ok_btn = QPushButton("OK") #OK
        self.cancel_btn = QPushButton("Cancel") #Cancel
        
        layout = QGridLayout()
        
        notice_utilization = """
        Dedicated tool for image segmentation.
        Import two elements:
        - A compressed file in zip format including only one or more images
        - A file in h5 format which is an image segmentation model
        - A file in txt format which is located classes names (optional)
        """
        
        layout.addWidget(QLabel(notice_utilization), 0, 1)        
        layout.addWidget(QLabel("Image:"), 1, 0)        
        layout.addWidget(self.filename_edit_image, 1, 1)
        layout.addWidget(self.file_name_image, 1, 2)
        layout.addWidget(QLabel("Model:"), 2, 0)        
        layout.addWidget(self.filename_edit_model, 2, 1)
        layout.addWidget(self.file_name_model, 2, 2)
        layout.addWidget(QLabel("Class (optional):"), 3, 0)        
        layout.addWidget(self.filename_edit_class_name, 3, 1)
        layout.addWidget(self.file_name_class_name, 3, 2)
        
        layout.addWidget(self.ok_btn, 4, 1)
        layout.addWidget(self.cancel_btn, 4, 2)
        self.setLayout(layout)
               
        self.ok_btn.clicked.connect(self.accept) #OK
        self.cancel_btn.clicked.connect(self.reject) #Cancel
        
        self.setFixedHeight(400)
        
    def open_file_dialog_image(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Images (*.png *.jpg *.zip *bmp *webp *tif *tiff)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_image.setText(str(path))     
            name_file = str(path)
            if name_file.endswith((".jpg",".JPG",".png",".PNG",".tiff",".tif",".bmp",".webp")):
                print('Image DETECTED')
            elif name_file.endswith((".zip",".ZIP")):
                print('Compressed file DETECTED')
            else:
                print('Image NOT DETECTED')
            
    def open_file_dialog_model(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Model (*.ilp *.h5)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_model.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".ilp"):
                print('Ilastik model DETECTED')
            elif name_file.endswith(".h5"):
                print('Tensorflow model DETECTED')
            else:
                print('Model NOT DETECTED')

    def open_file_dialog_class_name(self):
        filename, _ = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "File (*.txt)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_class_name.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".txt"):
                print('.txt file DETECTED')
            elif len(name_file)==0:
                print('No .txt file DETECTED')
            else:
                print('.txt file NOT DETECTED')
                                              
class Run_interface_segmentation:
    def __init__(self,x,y_list,current_viewer,output_dir):
        self.idx = x #
        self.image_zip = y_list[0] # fichier compresser en zip
        self.model = y_list[1] # fichier .h5 ou .ilp
        self.classe_file = y_list[2] # fichier txt
        self.napari_current_viewer = current_viewer
        self.temp_output_dir = output_dir
        self.nbr_classe = None
        self.label_name_current = None
        
    def run_model(self):
        """
        Connaitre le type de modele : ilastik ou tensorflow

        Returns:
            tuple: ('Image segmentation', type de modele)
        """
        dico = {1:'Image Segmentation',2:'Object Classification',3:'Image Classification',4:'Detection'}
        if self.model.endswith(".ilp"):
            root_pc = str(pathlib.Path.home()).split("\\")[0]+"\\Program Files"
            check_version = [ix for ix in os.listdir(root_pc) if ix.find('ilastik')!=-1]
            if len(check_version)==0:
                show_info('ILASTIK NOT INSTALLED')
                return (dico[self.idx],'') #NO ILASTIK INSTALLED
            else:
                ilastik_version = check_version[0]
                show_info('ILASTIK VERSION:'+ilastik_version)
                return ('Image Segmentation',ilastik_version)
        #elif self.model.endswith(".h5"):
        #    return ('Image Segmentation','Run tensorflow')
        else:
            return ('Image Segmentation','Run tensorflow')
    
    def run_tensorflow_segmentation(self,napari_viewer,tensorflow_version,id_uq_seg): 
        import tensorflow as tf
        from tensorflow.keras import backend as K
        from skimage.transform import resize
        from skimage import img_as_bool  

        model_New = tf.keras.models.load_model(self.model,custom_objects={'dice_coefficient': dice_coefficient})
        _, IMG_HEIGHT, IMG_WIDTH,IMG_CHANNELS = list(model_New.input.shape) # Taille INPUT
        h_p,w_p = get_output_size_image(model_New) # Taille OUTPUT
        self.nbr_classe = list(model_New.output_shape)[-1] # Nombre de classe
 
        if self.image_zip.endswith(".zip"):
            tmp_file = tempfile.TemporaryDirectory()
            with ZipFile(self.image_zip,'r') as zipObject:
                listOfFileNames = zipObject.namelist()
                if len(listOfFileNames)==1:
                    zipObject.extract(listOfFileNames[0],path=tmp_file.name)
                    image_path = os.path.join(tmp_file.name,listOfFileNames[0])
                        
                    matrix_img = imread(image_path)
                    if len(matrix_img.shape)==2:
                        or_h,or_w  = matrix_img.shape
                    else:
                        or_h,or_w,_  = matrix_img.shape
                        
                    # RESIZE
                    img = resize(matrix_img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
                        
                    # INPUT
                    if len(matrix_img.shape)==2:
                        X_ensemble = np.zeros((1, IMG_HEIGHT, IMG_WIDTH), dtype=img.dtype)
                    else:
                        X_ensemble = np.zeros((1, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
                    X_ensemble[0] = img

                    # INFERENCE
                    preds_test = model_New.predict(X_ensemble, verbose=1)
                    proba_img = preds_test[0,...]
                        
                    # RESIZE
                    final_output = resize(proba_img, (or_h, or_w), mode='constant', preserve_range=True)
                        
                    # POST
                    ## THRESHOLD
                    if self.nbr_classe==1:
                        final_output = final_output[:,:,0]
                        seuil = np.mean(final_output)
                        final_output = final_output > seuil
                    ## ARGMAX BETWEEN CHANNEL
                    else:
                        preds_test_t = np.argmax(final_output,axis=2)                     
                        for j in range(self.nbr_classe):
                            final_output[:,:,j] = (preds_test_t==j)*j
                        final_output = final_output.sum(axis=2)
                        final_output = final_output.astype('int') 
                    print(np.unique(final_output.flatten()))
                    return [
                        napari_viewer.add_image(matrix_img,name="Image_"+str(id_uq_seg)),
                        napari_viewer.add_labels(final_output,name="mask_"+str(id_uq_seg))
                    ]
                else:
                    pbar = progress(range(len(listOfFileNames)))
                        
                    # INPUT
                    IMG_RGB_list = np.zeros((len(listOfFileNames),IMG_HEIGHT, IMG_WIDTH,IMG_CHANNELS))
                    SHAPE_h_list = []
                    SHAPE_w_list = [] 
                    tmp_rgb_image_original = []
                    for i in pbar:
                        image_name = listOfFileNames[i]
                        zipObject.extract(listOfFileNames[i],path=tmp_file.name)
                        image_path = os.path.join(tmp_file.name,listOfFileNames[i])

                        matrix_img = imread(image_path) #image_path : PNG, JPG GOOD
                        or_h,or_w,_  = matrix_img.shape
                        tmp_rgb_image_original.append(matrix_img)
                        # RESIZE
                        img = resize(matrix_img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
                        IMG_RGB_list[i,:,:,:] = img
                        SHAPE_h_list.append(or_h)
                        SHAPE_w_list.append(or_w)
                    
                    stack_image_rgb = np.zeros((len(listOfFileNames),np.max(SHAPE_h_list),np.max(SHAPE_w_list),IMG_CHANNELS), dtype=np.uint8)
                    for i in range(len(listOfFileNames)):
                        stack_image_rgb[i,...]=tmp_rgb_image_original[i]
                    
                    # INFERENCE
                    preds_test = model_New.predict(IMG_RGB_list, verbose=1)

                    PROB_list = np.zeros((len(listOfFileNames),np.max(SHAPE_h_list),np.max(SHAPE_w_list)))
                    for i in pbar:
                            
                        pr_shape = preds_test.shape
                        pr_shape_int = len(pr_shape)
                        or_h, or_w = SHAPE_h_list[i],SHAPE_w_list[i]
                        # RESIZE
                        final_output = resize(preds_test[i,...], (or_h, or_w), mode='constant', preserve_range=True)[:,:,0]
                        PROB_list[i,...] = final_output

                        pbar.set_description("sds")
                    # POST                                            
                    ## THRESHOLD
                    if self.nbr_classe==1:
                        seuil = np.mean(PROB_list)
                        stack_masks = PROB_list > seuil
                    ## ARGMAX BETWEEN CHANNEL
                    else:
                        n,_h,_w,_ = PROB_list.shape
                        stack_maskss = np.zeros((n,_h,_w))
                        for ix in range(n):
                            current_stack = PROB_list[ix,...]
                            preds_test_t = np.argmax(current_stack,axis=2)
                            for j in range(self.nbr_classe):
                                current_stack[:,:,j] = (preds_test_t==j)*j
                            maskss = current_stack.sum(axis=2)
                            stack_masks[ix,...] = maskss
                            
                    return [
                        napari_viewer.add_image(stack_image_rgb,name="Image_"+str(id_uq_seg)),
                        napari_viewer.add_labels(stack_masks,name="mask_"+str(id_uq_seg))
                    ]
        elif "https://" in self.image_zip or "http://" in self.image_zip:
            img = Image.open(urlopen(self.image_zip))
            matrix_img = np.array(img)
            or_h,or_w,_  = matrix_img.shape
                        
            # RESIZE
            img = resize(matrix_img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
                        
            # INPUT
            X_ensemble = np.zeros((1, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
            X_ensemble[0] = img

            # INFERENCE
            preds_test = model_New.predict(X_ensemble, verbose=1)
            proba_img = preds_test[0,...]
                        
            # RESIZE
            final_output = resize(proba_img, (or_h, or_w), mode='constant', preserve_range=True)
                        
            # POST
            ## THRESHOLD
            if self.nbr_classe==1:
                final_output = final_output[:,:,0]
                seuil = np.mean(final_output)
                final_output = final_output > seuil
            ## ARGMAX BETWEEN CHANNEL
            else:
                preds_test_t = np.argmax(final_output,axis=2)
                for j in range(self.nbr_classe):
                    final_output[:,:,j] = (preds_test_t==j)*j
                final_output = final_output.sum(axis=2)
            return [
                napari_viewer.add_image(matrix_img,name="Image_"+str(id_uq_seg)),
                napari_viewer.add_labels(final_output,name="mask_"+str(id_uq_seg))
            ]

####
####
####
####
####
####
####
    
    def run_ilastik(self,ilastik_version):
        
        with ZipFile(self.image_zip,'r') as zipObject:
            
            listOfFileNames = zipObject.namelist()
            print(listOfFileNames)
            # lazy_arrays = []
            # lazy_masks = []
            
            if len(listOfFileNames)==1:
                zipObject.extract(listOfFileNames[0],path=self.temp_output_dir.name)
                image_path = os.path.join(self.temp_output_dir.name,listOfFileNames[0])
                projet_path = '--project='+self.model
                path_to_run = os.path.join("C:/Program Files",os.path.join(ilastik_version,"ilastik.exe"))
                recevoir = '--output_filename_format="'+os.path.join(self.temp_output_dir.name,listOfFileNames[0][:-4])+'_mask.jpg"'
                subprocess.run([path_to_run,
                                    '--headless',
                                    projet_path,
                                    '--export_source=Simple Segmentation',
                                    '--raw_data="'+image_path+'"',
                                    recevoir])
                print("IMAGE: 1")
            else:
                for i in progress(range(len(listOfFileNames))):
                    image_name = listOfFileNames[i]
                    
                    image_folder = self.temp_output_dir.name+'\\'+image_name[:-4]
                    os.mkdir(image_folder)
                    zipObject.extract(listOfFileNames[i],path=image_folder)
                    image_path = os.path.join(image_folder,listOfFileNames[i])
                                        
                    projet_path = '--project='+self.model
                    path_to_run = os.path.join("C:/Program Files",os.path.join(ilastik_version,"ilastik.exe"))
                    recevoir = '--output_filename_format="'+os.path.join(image_folder,image_name[:-4])+'_mask.jpg"'
                    subprocess.run([path_to_run,
                                    '--headless',
                                    projet_path,
                                    '--export_source=Simple Segmentation',
                                    '--raw_data="'+image_path+'"',
                                    recevoir])
                print("IMAGE:",i+1)
                    
    def color_panel(self,id_uq_seg):

        path_to_txt = self.classe_file
        slf_nb_class = self.nbr_classe+1
        if len(path_to_txt)==0:
            print("NO CLASS LABEL REFERENCED")
            
            idn_classe = [str(i) if i!=0 else 'Brush' for i in range(slf_nb_class)]
        else:
            with open(path_to_txt) as file:
                lines = file.readlines()
            file.close()
            idn_classe = [x.split('\n')[0] for x in lines]
            if len(idn_classe)==self.nbr_classe:
                print("CLASS LABEL REFERENCED")
                pass
            else:
                print("TOTAL CLASS LABEL NOT EQUAL TO MODEL")
                idn_classe = [str(i) if i!=0 else 'Brush' for i in range(slf_nb_class)]
                
        self.label_name_current = "mask_"+str(id_uq_seg)
        current_layer_bbx = self.napari_current_viewer.layers[self.label_name_current]

        def open_name_classe(item):
            idx = idn_classe.index(item.text())
            current_layer_bbx.mode = "PAINT"
            current_layer_bbx.selected_label = int(idx)
            current_class_displayed=None

        list_class_select = QListWidget()
        print(idn_classe)
        print(self.nbr_classe)
        slf_nb_class = self.nbr_classe+1
        for n,idx in zip(idn_classe,range(slf_nb_class)):
            i = QListWidgetItem(n)
            col_temp = current_layer_bbx.get_color(idx)
            if idx==0:
                pass
            else:
                i.setBackground(QColor(int(col_temp[0]*255),int(col_temp[1]*255),int(col_temp[2]*255),int(col_temp[3]*255)))
            list_class_select.addItem(i)
                        
        list_class_select.currentItemChanged.connect(open_name_classe)
        self.napari_current_viewer.window.add_dock_widget([list_class_select], area='right',name="Images")
        list_class_select.setCurrentRow(0)
        

class Image_classification(QDialog):
    
    def __init__(self, parent: QWidget):
        print("Image classification OPEN")
        super().__init__(parent)
        self.setWindowTitle("Image classification")
        self.number = QSpinBox()
        
        self.filename_edit = QLineEdit()  
        
        self.file_name_image = QPushButton("File")
        self.filename_edit_image = QLineEdit()  
        self.file_name_image.clicked.connect(self.open_file_dialog_image)
        
        self.file_name_model = QPushButton("File")
        self.filename_edit_model = QLineEdit()  
        self.file_name_model.clicked.connect(self.open_file_dialog_model)

        self.file_name_class_name = QPushButton("File")
        self.filename_edit_class_name = QLineEdit()  
        self.file_name_class_name.clicked.connect(self.open_file_dialog_class_name)
        
        self.ok_btn = QPushButton("OK") #OK
        self.cancel_btn = QPushButton("Cancel") #Cancel
        
        layout = QGridLayout()
        # layout.addWidget(QLabel("Number:"), 0, 0)
        # layout.addWidget(self.number, 0, 1)
        
        notice_utilization = """
        Dedicated tool for image classification.
        Import three required elements:
        - A compressed file in zip format including only one or more images
        - A file in h5 format which is an image classification model
        - A file in txt format where we found a list of classes associated to the model
        """
        
        layout.addWidget(QLabel(notice_utilization), 0, 1)        
        layout.addWidget(QLabel("Image:"), 1, 0)        
        layout.addWidget(self.filename_edit_image, 1, 1)
        layout.addWidget(self.file_name_image, 1, 2)
        layout.addWidget(QLabel("Model:"), 2, 0)        
        layout.addWidget(self.filename_edit_model, 2, 1)
        layout.addWidget(self.file_name_model, 2, 2)
        layout.addWidget(QLabel("Class:"), 3, 0)        
        layout.addWidget(self.filename_edit_class_name, 3, 1)
        layout.addWidget(self.file_name_class_name, 3, 2)
        
        layout.addWidget(self.ok_btn, 4, 1)
        layout.addWidget(self.cancel_btn, 4, 2)
        self.setLayout(layout)
               
        self.ok_btn.clicked.connect(self.accept) #OK
        self.cancel_btn.clicked.connect(self.reject) #Cancel
        
        self.setFixedHeight(400)
        
    def open_file_dialog_image(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Images (*.png *.jpg *.zip *bmp *webp *tif *tiff)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_image.setText(str(path))     
            name_file = str(path)
            if name_file.endswith((".zip")):
                print('Compressed file DETECTED')
            else:
                print('Compressed file NOT DETECTED')
            
    def open_file_dialog_model(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Model (*.h5)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_model.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".ilp"):
                print('Ilastik model DETECTED')
            elif name_file.endswith(".h5"):
                print('Tensorflow model DETECTED')
            else:
                print('Model NOT DETECTED')

    def open_file_dialog_class_name(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "File (*.txt)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_class_name.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".txt"):
                print('.txt file DETECTED')
            else:
                print('.txt file NOT DETECTED')

class Run_interface_classification:
    def __init__(self,x,y_list,current_viewer,output_dir):
        self.idx = x
        self.image_zip = y_list[0]
        self.model = y_list[1]
        self.classe_file = y_list[2]
        self.napari_current_viewer = current_viewer
        self.temp_output_dir = output_dir
    def run_model(self):
        dico = {1:'Image Segmentation',2:'Object Classification',3:'Image Classification',4:'Detection'}
        if self.model.endswith(".ilp"):
            root_pc = str(pathlib.Path.home()).split("\\")[0]+"\\Program Files"
            check_version = [ix for ix in os.listdir(root_pc) if ix.find('ilastik')!=-1]
            if len(check_version)==0:
                show_info('ILASTIK NOT INSTALLED')
                return (dico[self.idx],'') #NO ILASTIK INSTALLED
            else:
                ilastik_version = check_version[0]
                show_info('ILASTIK VERSION:'+ilastik_version)
                return (dico[self.idx],ilastik_version)
        elif self.model.endswith(".h5"):
            return (dico[self.idx],'Run tensorflow')
    
    def run_tensorflow_classification(self,tensorflow_version):
        import tensorflow as tf
        from tensorflow.keras import backend as K
        import cv2
        import random
        
        model_New_load = tf.keras.models.load_model(self.model)
        
        _, IMG_SIZE_H, IMG_SIZE_W,_ = list(model_New_load.input.shape)
        
        def create_data(path):
            image_name = []
            data = []
            A = os.listdir(path)
            n = len(A)
            # X = np.zeros((n,IMG_SIZE_H, IMG_SIZE_W,3))  # An Array for images
            list_img = []
            for i,img in enumerate(A):  # iterate over each image per plants and weeds
                img_procee = img.lower()
                if img_procee.endswith(('.tif','.png','.jpg')):
                    image_name.append(img)
                    # img_array = cv2.imread(os.path.join(path,img))  # convert to array 
                    imgtp = Image.open(os.path.join(path,img))
                    img_array = np.array(imgtp)
                    list_img.append(img_array)
                    # new_array = cv2.resize(img_array, (IMG_SIZE_H, IMG_SIZE_W))  # resize to normalize data size
                    # X[i,:,:,:] = new_array
            # X = np.array(X).reshape(-1, IMG_SIZE_H, IMG_SIZE_W, 3)  # Reshape data in a form that is suitable for keras
            # return X,image_name
            return list_img,image_name
        
        if self.image_zip.endswith(".zip"):
            with ZipFile(self.image_zip,'r') as zipObject:
                listOfFileNames = zipObject.namelist()
                
                # Extraction des images dans un dossier temporaires
                pbar = progress(range(len(listOfFileNames)))
                for i in pbar:
                    image_name = listOfFileNames[i]
                    zipObject.extract(listOfFileNames[i],path=self.temp_output_dir.name)
                
                # Affecter les classes dans variable lines
                txt_file_path = self.classe_file
                with open(txt_file_path) as file:
                    lines = file.readlines()
                file.close()
                self.LABEL_CATEGORY = [x.split('\n')[0] for x in lines]
                
                # Affecter les images + noms dans variables
                list_img,image_name = create_data(self.temp_output_dir.name)
                image_name_temp = [ ix.split(".")[0] for ix in image_name]
                
                # Preprocessing
                n = len(image_name_temp) # nombre dimages
                X = np.zeros((n,IMG_SIZE_H, IMG_SIZE_W,3))  # An Array for images
                for ix in list_img:
                    new_array = cv2.resize(ix, (IMG_SIZE_H, IMG_SIZE_W))  # resize to normalize data size
                    X[i,:,:,:] = new_array
                X_test = X.astype('float32')/255.
                
                # Processing
                dico_pred = {} # dico pred result
                n = len(image_name_temp) # nombre dimages
                m = len(self.LABEL_CATEGORY) # nombre de classes
                self.prob_class = [] # Probabilite max
                res_classe = []
                            
                preds_test_ = model_New_load.predict(X_test, verbose=1)
                for ix in range(n):
                    prod_res = preds_test_[ix,...]
                    id_class = np.argmax(prod_res)
                    self.prob_class.append(np.max(prod_res))
                    res_classe.append(self.LABEL_CATEGORY[id_class])
                dico_pred['nom']=image_name_temp
                dico_pred['prediction']=res_classe
                dico_pred['prob']=self.prob_class

                self.dico_output_prediction = dico_pred
        else:
            head,image_name = os.path.split(self.image_zip)
            if "https://" in self.image_zip or "http://" in self.image_zip:                
                img = Image.open(urlopen(self.image_zip))
                img.save(os.path.join(self.temp_output_dir.name,"internet.jpg")) 
                img_array = np.array(img)
                
                # Affecter les classes dans variable lines
                txt_file_path = self.classe_file
                with open(txt_file_path) as file:
                    lines = file.readlines()
                file.close()
                self.LABEL_CATEGORY = [x.split('\n')[0] for x in lines]
                
                # Preprocessing
                X = np.zeros((1,IMG_SIZE_H, IMG_SIZE_W,3))  # An Array for images
                new_array = cv2.resize(img_array, (IMG_SIZE_H, IMG_SIZE_W))  # resize to normalize data size
                X[0,:,:,:] = new_array
                
                X_test = X.astype('float32')/255.
                
                # Processing
                dico_pred = {} # dico pred result
                            
                preds_test_ = model_New_load.predict(X_test, verbose=1)
                
                prod_res = preds_test_[0,...]
                id_class = np.argmax(prod_res)
                self.prob_class = [np.max(prod_res)] # Probabilite max
                res_classe = [self.LABEL_CATEGORY[id_class]]
                
                dico_pred['nom']=[image_name]
                dico_pred['prediction']=res_classe
                dico_pred['prob']=self.prob_class

                self.dico_output_prediction = dico_pred
    
    def image_vis(self):

        def convertTableToDataFrame(self_table):
            rowCount = self_table.rowCount()
            columnCount = 3
                
            Nom_image = []
            Prediction = []
            for row in range(rowCount):
                rowData = []
                for column in range(columnCount):
                    widgetItem = self_table.item(row, column)
                    if widgetItem and widgetItem.text:
                        rowData.append(widgetItem.text())
                    else:
                        widget = self_table.cellWidget(row, column)
                        if isinstance(widget,QComboBox):
                            rowData.append(widget.currentText())
                Nom_image.append(rowData[0])
                Prediction.append(rowData[1])
                
                dico_data_displayed = {"nom":Nom_image,"prediction":Prediction}

            for cle in self.dico_output_prediction:
                if cle not in ['nom','prediction','prob']:
                    dico_data_displayed[cle] = self.dico_output_prediction[cle]

            df = pd.DataFrame(dico_data_displayed)
            return df


        def table_to_widget(table: dict, LABEL_CATEGORY: list) -> QWidget:
            """
            Takes a table given as dictionary with strings as keys and numeric arrays as values and returns a QWidget which
            contains a QTableWidget with that data.
            """
            save_button = QPushButton("Save as csv...")
            
            @save_button.clicked.connect
            def save_trigger():
                filename, _ = QFileDialog.getSaveFileName(save_button, "Save as csv...", ".", "*.csv")
                df = convertTableToDataFrame(self.v_table)
                df.to_csv(filename)    
                show_info(f"{filename} SAVED")        
                
            widget = QWidget()
            widget.setWindowTitle("Prediction")
            widget.setLayout(QVBoxLayout())
            self.v_table = TableWidget(table,self.LABEL_CATEGORY)
            widget.setFixedWidth(500)
            widget.setFixedHeight(200)
            widget.layout().addWidget(self.v_table)
            widget.layout().addWidget(save_button)
            return widget

        names = []
        dico_name = {}
        for image_name in os.listdir(self.temp_output_dir.name):
            image_name_without_extension, _ = os.path.splitext(image_name)
            dico_name[image_name_without_extension]=image_name
            names.append(image_name_without_extension)

        def open_name(item):
            name = item.text()
            image_name_without_extension, extension_format = os.path.splitext(dico_name[name])
            fname = os.path.join(self.temp_output_dir.name,dico_name[name]) 
            self.napari_current_viewer.layers.select_all()
            self.napari_current_viewer.layers.remove_selected()  
            data_segment0 = imread(fname)
            self.napari_current_viewer.add_image(data_segment0,name=f'{image_name_without_extension}')
        
        widget = QWidget()
        list_widget = QListWidget()
        for n in names:
            list_widget.addItem(n)    
        list_widget.currentItemChanged.connect(open_name)
        self.napari_current_viewer.window.add_dock_widget([list_widget], area='right',name="Images")
        list_widget.setCurrentRow(0)

        dock_widget = table_to_widget(self.dico_output_prediction,self.LABEL_CATEGORY)
        self.napari_current_viewer.window.add_dock_widget(dock_widget, area='right') 
        return self.dico_output_prediction



class Object_detection(QDialog):
    
    def __init__(self, parent: QWidget):
        print("Object detection OPEN")
        super().__init__(parent)
        self.setWindowTitle("Object detection")
        self.number = QSpinBox()
        
        self.filename_edit = QLineEdit()  
        
        self.file_name_image = QPushButton("File")
        self.filename_edit_image = QLineEdit()  
        self.file_name_image.clicked.connect(self.open_file_dialog_image)
        
        self.file_name_model = QPushButton("File")
        self.filename_edit_model = QLineEdit()  
        self.file_name_model.clicked.connect(self.open_file_dialog_model)

        self.file_name_class_name = QPushButton("File")
        self.filename_edit_class_name = QLineEdit()  
        self.file_name_class_name.clicked.connect(self.open_file_dialog_class_name)
        
        self.ok_btn = QPushButton("OK") #OK
        self.cancel_btn = QPushButton("Cancel") #Cancel
        
        layout = QGridLayout()
        
        notice_utilization = """
        Dedicated tool for object detection.
        Import two elements:
        - A compressed file in zip format including only one or more images
        - A file in h5 format which is an object detection model
        - A file in txt format which is located classes names (optional)
        """
        
        layout.addWidget(QLabel(notice_utilization), 0, 1)        
        layout.addWidget(QLabel("Image:"), 1, 0)        
        layout.addWidget(self.filename_edit_image, 1, 1)
        layout.addWidget(self.file_name_image, 1, 2)
        layout.addWidget(QLabel("Model:"), 2, 0)        
        layout.addWidget(self.filename_edit_model, 2, 1)
        layout.addWidget(self.file_name_model, 2, 2)
        layout.addWidget(QLabel("Class (optional):"), 3, 0)        
        layout.addWidget(self.filename_edit_class_name, 3, 1)
        layout.addWidget(self.file_name_class_name, 3, 2)
        
        layout.addWidget(self.ok_btn, 4, 1)
        layout.addWidget(self.cancel_btn, 4, 2)
        self.setLayout(layout)
               
        self.ok_btn.clicked.connect(self.accept) #OK
        self.cancel_btn.clicked.connect(self.reject) #Cancel
        
        self.setFixedHeight(400)
        
    def open_file_dialog_image(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Images (*.png *.jpg *.zip *bmp *webp *tif *tiff)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_image.setText(str(path))     
            name_file = str(path)
            if name_file.endswith((".jpg",".JPG",".png",".PNG",".tiff",".tif",".bmp",".webp")):
                print('Image DETECTED')
            elif name_file.endswith((".zip",".ZIP")):
                print('Compressed file DETECTED')
            else:
                print('Image NOT DETECTED')
            
    def open_file_dialog_model(self):
        filename, ok = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "Model (*.ilp *.h5)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_model.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".ilp"):
                print('Ilastik model DETECTED')
            elif name_file.endswith(".h5"):
                print('Tensorflow model DETECTED')
            else:
                print('Model NOT DETECTED')

    def open_file_dialog_class_name(self):
        filename, _ = QFileDialog.getOpenFileName(
            self,
            "Select a File", 
            str(pathlib.Path.home()), 
            "File (*.txt)"
        )
        if filename:
            path = pathlib.Path(filename)
            self.filename_edit_class_name.setText(str(path))  
            name_file = str(path)
            if name_file.endswith(".txt"):
                print('.txt file DETECTED')
            elif len(name_file)==0:
                print('No .txt file DETECTED')
            else:
                print('.txt file NOT DETECTED')
                
class Run_interface_detection:
    def __init__(self,x,y_list,current_viewer,output_dir):
        self.idx = x #
        self.image_zip = y_list[0] # fichier compresser en zip
        self.model = y_list[1] # fichier .h5 ou .ilp
        self.classe_file = y_list[2] # fichier txt       
        self.napari_current_viewer = current_viewer
        self.temp_output_dir = output_dir 
        self.nbr_classe = None
        self.label_name_current = None
        
    def run_tensorflow_detection(self,napari_viewer,id_uq_detect):
        import tensorflow as tf
        import cv2
        import colorsys
        import random
       
        model_New = tf.keras.models.load_model(self.model,compile=False)
        _, IMG_HEIGHT, IMG_WIDTH,IMG_CHANNELS = list(model_New.input.shape) # Taille INPUT
        self.nbr_classe = list(model_New.output_shape)[-1] # Nombre de classe
        
        # Affecter les classes dans variable lines
        txt_file_path = self.classe_file
        with open(txt_file_path) as file:
            lines = file.readlines()
        file.close()
        self.LABEL_CATEGORY = [x.split('\n')[0] for x in lines]
        
        self.num_classes = len(self.LABEL_CATEGORY)
        classes_list = {i:self.LABEL_CATEGORY[i] for i in range(self.num_classes)}
        hsv_tuples = [(1.0 * x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
        colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

        # Importer les images
        if self.image_zip.endswith('.zip'):
            with ZipFile(self.image_zip,'r') as zipObject:
                listOfFileNames = zipObject.namelist()
                n = len(listOfFileNames)
                
                images_data_or = []
                images_data_rs = []
                
                results_mess = []
                results_coords = []
                results_color = []
                SHAPE_h_list = []
                SHAPE_w_list = []
                for i in range(n):
                    print(i)
                    zipObject.extract(listOfFileNames[i],path=self.temp_output_dir.name)
                    image_path = os.path.join(self.temp_output_dir.name,listOfFileNames[i])
                    original_image = cv2.imread(image_path)
                    original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
                    h_,w_,_ = original_image.shape
                    SHAPE_h_list.append(h_)
                    SHAPE_w_list.append(w_)
                    images_data_or.append(original_image)
                    #PREPROCESSING
                    img_rsz = cv2.resize(original_image,(IMG_HEIGHT, IMG_WIDTH))
                    img_rsz = img_rsz / 255
                    images_data_rs = [img_rsz]
                    images_data_rs = np.asarray(images_data_rs).astype(np.float32)
                
                    #PROCESSING
                    ##DARKNET
                    batch = tf.constant(images_data_rs)
                    pred_bbox_ = model_New(batch)
                
                    ##NON_MAX_SUPPRESSION
                    boxes = pred_bbox_[:, :, 0:4]
                    pred_conf = pred_bbox_[:, :, 4:]
                    
                    boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
                        boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                        scores=tf.reshape(
                            pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                        max_output_size_per_class=50,
                        max_total_size=50,
                        iou_threshold=0.45,
                        score_threshold=0.25
                    )

                    pred_bboxes = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]

                    random.seed(0)
                    random.shuffle(self.colors)
                    random.seed(None)
                    bbox_rect = []
                    
                    image_h, image_w, _ = original_image.shape
                    out_boxes, out_scores, out_classes, num_boxes = pred_bboxes
                    for k in range(num_boxes[0]):
                        if int(out_classes[0][k]) < 0 or int(out_classes[0][k]) > self.num_classes: continue
                        coor = out_boxes[0][k]
                        coor[0] = int(coor[0] * image_h) #y1
                        coor[2] = int(coor[2] * image_h) #y2
                        coor[1] = int(coor[1] * image_w) #x1
                        coor[3] = int(coor[3] * image_w) #x2
                        x1 = coor[1]
                        x2 = coor[3]
                        y1 = coor[0]
                        y2 = coor[2]
                        
                        # class BBOX
                        class_ind = int(out_classes[0][k])
                        bbox_mess = classes_list[class_ind]
                        results_mess.append(bbox_mess)
                        
                        # color BBOX
                        class_ind = int(out_classes[0][k])
                        bbox_color_ = self.colors[class_ind]
                        bbox_color = [bbox_color_[0]/255,bbox_color_[1]/255,bbox_color_[2]/255]
                        results_color.append(bbox_color)
                        
                        # coords BBOX
                        bbox_rect.append(np.array([[i,y1, x1], [i,y2, x1], [i,y2, x2], [i,y1, x2]])) # coords BBOX
                    results_coords+=bbox_rect
        
        else:
            head,image_name = os.path.split(self.image_zip)
            if "https://" in self.image_zip or "http://" in self.image_zip:                
                img = Image.open(urlopen(self.image_zip))
                img.save(os.path.join(self.temp_output_dir.name,"internet.jpg")) 
                original_image = np.array(img)
                i=0
                images_data_or = []
                images_data_rs = []
                
                results_mess = []
                results_coords = []
                results_color = []
                SHAPE_h_list = []
                SHAPE_w_list = []

                #original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
                h_,w_,_ = original_image.shape
                SHAPE_h_list.append(h_)
                SHAPE_w_list.append(w_)
                images_data_or.append(original_image)
                #PREPROCESSING
                img_rsz = cv2.resize(original_image,(IMG_HEIGHT, IMG_WIDTH))
                img_rsz = img_rsz / 255
                images_data_rs = [img_rsz]
                images_data_rs = np.asarray(images_data_rs).astype(np.float32)
                
                #PROCESSING
                ##DARKNET
                batch = tf.constant(images_data_rs)
                pred_bbox_ = model_New(batch)
                
                ##NON_MAX_SUPPRESSION
                boxes = pred_bbox_[:, :, 0:4]
                pred_conf = pred_bbox_[:, :, 4:]
                    
                boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
                    boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                    scores=tf.reshape(
                        pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                    max_output_size_per_class=50,
                    max_total_size=50,
                    iou_threshold=0.45,
                    score_threshold=0.25
                )

                pred_bboxes = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]

                random.seed(0)
                random.shuffle(self.colors)
                random.seed(None)
                bbox_rect = []
                    
                image_h, image_w, _ = original_image.shape
                out_boxes, out_scores, out_classes, num_boxes = pred_bboxes
                for k in range(num_boxes[0]):
                    if int(out_classes[0][k]) < 0 or int(out_classes[0][k]) > self.num_classes: continue
                    coor = out_boxes[0][k]
                    coor[0] = int(coor[0] * image_h) #y1
                    coor[2] = int(coor[2] * image_h) #y2
                    coor[1] = int(coor[1] * image_w) #x1
                    coor[3] = int(coor[3] * image_w) #x2
                    x1 = coor[1]
                    x2 = coor[3]
                    y1 = coor[0]
                    y2 = coor[2]
                        
                    # class BBOX
                    class_ind = int(out_classes[0][k])
                    bbox_mess = classes_list[class_ind]
                    results_mess.append(bbox_mess)
                        
                    # color BBOX
                    class_ind = int(out_classes[0][k])
                    bbox_color_ = self.colors[class_ind]
                    bbox_color = [bbox_color_[0]/255,bbox_color_[1]/255,bbox_color_[2]/255]
                    results_color.append(bbox_color)
                        
                    # coords BBOX
                    bbox_rect.append(np.array([[i,y1, x1], [i,y2, x1], [i,y2, x2], [i,y1, x2]])) # coords BBOX
                results_coords+=bbox_rect
                listOfFileNames = [1]
        
                
        properties = {
            'label': results_mess,
        }
        
        text_parameters = {
            'string': '{label}',
            'size': 12,
            'color': 'red',
            'anchor': 'upper_left',
            'translation': [-3, 0]
        }

        stack_image_rgb = np.zeros((len(listOfFileNames),np.max(SHAPE_h_list),np.max(SHAPE_w_list),IMG_CHANNELS), dtype=np.uint8)
        for i in range(len(listOfFileNames)): 
            stack_image_rgb[i,...][:images_data_or[i].shape[0], :images_data_or[i].shape[1]] = images_data_or[i]
        images_layer = self.napari_current_viewer.add_image(stack_image_rgb,name="Image_"+str(id_uq_detect))

        shapes_layer = self.napari_current_viewer.add_shapes(
            results_coords,
            edge_color=results_color,
            face_color='transparent', edge_width=5,
            properties=properties,
            text=text_parameters,
            name='Result_'+str(id_uq_detect),
        )
        
###
        # current_layer_bbx = self.napari_current_viewer.layers['Result_'+str(id_uq_detect)]
            
        # def open_name_classe(item):
        #     idx = self.LABEL_CATEGORY.index(item.text())
        #     current_layer_bbx.mode = 'add_rectangle' 
        #     current_layer_bbx.current_edge_color = (colors[idx][0]/255,colors[idx][1]/255,colors[idx][2]/255)
        #     current_class_displayed=None

        # list_class_select = QListWidget()
        # for n,idx in zip(self.LABEL_CATEGORY,range(num_classes)):
        #     i = QListWidgetItem(n)
        #     i.setBackground(QColor(colors[idx][0]/255,colors[idx][1]/255,colors[idx][2]/255))
        #     list_class_select.addItem(i)
                        
        # list_class_select.currentItemChanged.connect(open_name_classe)
        # self.napari_current_viewer.window.add_dock_widget([list_class_select], area='right',name="Images")
        # list_class_select.setCurrentRow(0)

###
        
        print('END',colors)
        return [
                images_layer,
                shapes_layer  
                ]
    # def color_panel(self,id_uq_seg):
    #     print("==")
          
    def color_panel(self,id_uq_seg):
        current_layer_bbx = self.napari_current_viewer.layers['Result_'+str(id_uq_detect)]
            
        def open_name_classe(item):
            idx = self.LABEL_CATEGORY.index(item.text())
            current_layer_bbx.mode = 'add_rectangle' 
            # current_layer_bbx.text = item.text()
            print(current_layer_bbx.properties['label'])
            # klff = list(current_layer_bbx.properties['label'])
            # klff.append(item.text())
            # current_layer_bbx.properties['label'] = np.array(klff,dtype=object)
            current_layer_bbx.current_edge_color = (self.colors[idx][0]/255,self.colors[idx][1]/255,self.colors[idx][2]/255)
            current_class_displayed=None

        list_class_select = QListWidget()
        for n,idx in zip(self.LABEL_CATEGORY,range(self.num_classes)):
            i = QListWidgetItem(n)
            i.setBackground(QColor(self.colors[idx][0],self.colors[idx][1],self.colors[idx][2]))
            list_class_select.addItem(i)
                        
        list_class_select.currentItemChanged.connect(open_name_classe)
        self.napari_current_viewer.window.add_dock_widget([list_class_select], area='right',name="Images")
        list_class_select.setCurrentRow(0)

id_uq_seg = 0          
id_uq_detect = 0
class ManiniWidget(QWidget):
    # your QWidget.__init__ can optionally request the napari viewer instance
    # in one of two ways:
    # 1. use a parameter called `napari_viewer`, as done here
    # 2. use a type annotation of 'napari.viewer.Viewer' for any parameter
    def __init__(self, napari_viewer):
        super().__init__()
        self.viewer = napari_viewer
        self.ix = None
        self.iy = None
        self.old_subfolder_image = None
        self.view_table = None
        self.v_table = None
        self.dico_output_prediction = None
        
        self.NOM_IMAGE_list_displayed = None
        self.SIZE_list_displayed = None
        self.CLASS_list_displayed = None
        self.DICO_DETAIL_list_displayed = None
        self.classe_dico_idx = None
        self.dico = None
        self.dico_name_extension = None

        
        self.output_directory = tempfile.TemporaryDirectory()
                
        click1 = QPushButton("Image segmentation")
        click1.clicked.connect(self._on_click1)           
        click_run = QPushButton("Run")
        click_run.clicked.connect(self._on_click_run)    
        click3 = QPushButton("Image classification")
        click3.clicked.connect(self._on_click3)    
        # click_save = QPushButton("Save")
        # click_save.clicked.connect(self._on_click_save)
        click4 = QPushButton("Object Detection")
        click4.clicked.connect(self._on_click4)
                
        self.setLayout(QGridLayout())       
    
        self.layout().addWidget(click1,1,0)
        self.layout().addWidget(click3,2,0)
        self.layout().addWidget(click4,3,0)
        
        self.layout().addWidget(click_run,4,0)
        # self.layout().addWidget(click_save,4,1)
        
    def _on_click1(self):
        dialog = Image_segmentation(self)  
        dialog.exec_()        
        print("Image segmentation CLOSE")
        if dialog.result() == QDialog.Accepted:
            self._on_click_idx = 1
            self._on_click_list = [dialog.filename_edit_image.text(),dialog.filename_edit_model.text(),dialog.filename_edit_class_name.text()]
            # self._on_click_list contient : chemin fichier .zip, chemin model, chemin fichier txt

    def _on_click3(self):
        dialog = Image_classification(self)  
        dialog.exec_()
        print("Image classification CLOSE")
        if dialog.result() == QDialog.Accepted:
            self._on_click_idx = 3
            self._on_click_list = [dialog.filename_edit_image.text(),dialog.filename_edit_model.text(),dialog.filename_edit_class_name.text()]
        
    def _on_click4(self):
        dialog = Object_detection(self)  
        dialog.exec_()        
        print("Object detection CLOSE")
        if dialog.result() == QDialog.Accepted:
            self._on_click_idx = 4
            self._on_click_list = [dialog.filename_edit_image.text(),dialog.filename_edit_model.text(),dialog.filename_edit_class_name.text()]
            
    def _on_click_run(self):
        global id_uq_seg
        global id_uq_detect
        if self._on_click_idx==1:
            id_uq_seg+=1
            a = Run_interface_segmentation(self._on_click_idx,self._on_click_list,self.viewer,self.output_directory)
            val1,val2 = a.run_model() # {1:'Image Segmentation',2:'Object Classification',3:'Image Classification',4:'Detection'}
            
            t1_start = process_time() #time start      
            if val2 != 'Run tensorflow':
                a.run_ilastik(val2) # Run ILASTIK
            else:  
                a.run_tensorflow_segmentation(self.viewer,val2,id_uq_seg) # Run TENSORFLOW
            t1_stop = process_time() #time stop
            print("TOTAL PROCESSING TIME:",np.round(t1_stop-t1_start,2),"seconds") 
            a.color_panel(id_uq_seg)
        elif self._on_click_idx==3:
            self.output_directory = tempfile.TemporaryDirectory()
            a = Run_interface_classification(self._on_click_idx,self._on_click_list,self.viewer,self.output_directory)
            val1,val2 = a.run_model() # {1:'Image Segmentation',2:'Object Classification',3:'Image Classification',4:'Detection'}
            t1_start = process_time()
            a.run_tensorflow_classification(val2)
            t1_stop = process_time() #time stop
            print("TOTAL PROCESSING TIME:",np.round(t1_stop-t1_start,2),"seconds")
            # self.v_table,self.dico_output_prediction = a.image_vis()
            self.dico_output_prediction = a.image_vis()
        elif self._on_click_idx==4:
            id_uq_detect+=1
            self.output_directory = tempfile.TemporaryDirectory()
            a = Run_interface_detection(self._on_click_idx,self._on_click_list,self.viewer,self.output_directory)
            # val1,val2 = a.run_model() # {1:'Image Segmentation',2:'Object Classification',3:'Image Classification',4:'Detection'}
            t1_start = process_time()
            a.run_tensorflow_detection(self.viewer,id_uq_detect)
            t1_stop = process_time() #time stop
            print("TOTAL PROCESSING TIME:",np.round(t1_stop-t1_start,2),"seconds")
            a.color_panel(id_uq_seg)
            
            
@magic_factory(call_button="save zip",layout="vertical")
def save_as_zip(layer_: Layer,layer_RGB : ImageData):
# def save_as_zip(layer_mask: LabelsData,layer_RGB : ImageData):
    # if str(type(layer_).__name__) not in ['Shapes','Labels']:
    #     show_error('Please choose Shapes or Labels')
    #     assert str(type(layer_).__name__) in ['Shapes','Labels'],'Please choose Shapes or Labels'
    # assert str(type(layer_).__name__) in ['Shapes','Labels'],'Please choose Shapes or Labels'
    layer_mask = layer_.data
    if str(type(layer_).__name__) == 'Labels':
        save_labels(layer_mask,layer_RGB)
    elif str(type(layer_).__name__) == 'Shapes':
        layer_class = layer_.properties['label']
        save_shape(layer_mask,layer_class,layer_RGB)
    else:
        show_error('Please choose Shapes or Labels')
        
def save_shape(layer_bbx,layer_class,layer_RGB):
    zip_dir = tempfile.TemporaryDirectory()
    save_button = QPushButton("Save as zip")
    filename, _ = QFileDialog.getSaveFileName(save_button, "Save as zip", ".", "zip")

    nbr_image = layer_RGB.shape 
    
    if len(nbr_image)==4:
        
        total_RGB = nbr_image[0]
        img_array = np.array(layer_RGB)
        # bbx_array = np.array(layer_bbx,dtype='int32')
        
        A = []
        for ix,iy in zip(layer_bbx,layer_class): 
            vg = np.zeros((4,4),dtype='object')
            vg[:,:3] = ix
            vg[:,3] = iy
            A.append(vg)
        # bbx_array = np.array(A,dtype='int32')
        bbx_array = np.array(A)
        
        for ix in tqdm(range(total_RGB),"Extracting"):
            data_RGB = img_array[ix,...]

            #recherche des bbx de image courant
            bbx_current = bbx_array[bbx_array[:,:,0]==[ix,ix,ix,ix]]
            # bbx_current = bbx_array[bbx_array[:,0]==[ix,ix,ix,ix]] 
            n = bbx_current.shape[0]
            if n==0:
                CX_ = ['']
                CY_ = ['']
                H_ = ['']
                W_ = ['']     
                C_ = ['']       
            else:
                CX_ = []
                CY_ = []
                H_ = []
                W_ = []
                C_ = []
                for i,j in zip(range(0,n,4),range(4,n+4,4)):
                    bbx_coord = bbx_current[i:j,:] 
                    _ , minr, minc, class_bbx = bbx_coord[0] # <-- adapter avec orientation du rectangle
                    _ , maxr, minc, _ = bbx_coord[1]
                    _ , maxr, maxc, _ = bbx_coord[2]
                    _ , minr, maxc, _ = bbx_coord[3]
                    cx = (int(minc)+int(maxc))/2
                    cy = (int(minr)+int(maxr))/2
                    h = int(maxr)-int(minr)
                    w = int(maxc)-int(minc)
                    CX_.append(cx)
                    CY_.append(cy)
                    H_.append(h)
                    W_.append(w)
                    C_.append(class_bbx)
            df = pd.DataFrame({'x':CX_,'y':CY_,'h':H_,'w':W_,'class':C_})
            df.to_csv(os.path.join(zip_dir.name,str(ix)+'.csv'),index=False)
            
            im_RGB = Image.fromarray(data_RGB)
            im_RGB1 = im_RGB.save(os.path.join(zip_dir.name,'RGB_'+str(ix)+".png"))
        
    elif len(nbr_image)==3:
        total_RGB = nbr_image[0]
        img_array = np.array(layer_RGB)
        bbx_array = np.array(layer_bbx,dtype='int32')

        #recherche des bbx de image courant
        n = bbx_array.shape[0]
        if n==0:
            CX_ = ['']
            CY_ = ['']
            H_ = ['']
            W_ = ['']        
        else:
            CX_ = []
            CY_ = []
            H_ = []
            W_ = []
            for ix in range(n):
                bbx_coord = bbx_array[ix]
                minr, minc = bbx_coord[0]
                maxr, minc = bbx_coord[1]
                maxr, maxc = bbx_coord[2]
                minr, maxc = bbx_coord[3]
                cx = (minc+maxc)/2
                cy = (minr+maxr)/2
                h = maxr-minr
                w = maxc-minc
                CX_.append(cx)
                CY_.append(cy)
                H_.append(h)
                W_.append(w)
        df = pd.DataFrame({'x':CX_,'y':CY_,'h':H_,'w':W_})
        df.to_csv(os.path.join(zip_dir.name,'1.csv'),index=False)
            
        im_RGB = Image.fromarray(data_RGB)
        im_RGB1 = im_RGB.save(os.path.join(zip_dir.name,"1.png"))
    shutil.make_archive(filename, 'zip', zip_dir.name)
    
    for ix in os.listdir(zip_dir.name):
        os.remove(os.path.join(zip_dir.name,ix))
    show_info('Compressed file done')

def save_labels(layer_mask,layer_RGB):
    zip_dir = tempfile.TemporaryDirectory()
    save_button = QPushButton("Save as zip")
    filename, _ = QFileDialog.getSaveFileName(save_button, "Save as zip", ".", "zip")

    nbr_image = layer_RGB.shape
    nbr_mask = layer_mask.shape
    
    if len(nbr_image)==4:
        assert nbr_image[0]==nbr_mask[0], "MASK AND RGB SIZE NOT EQUAL"
        
        total_RGB = nbr_image[0]
        img_array = np.array(layer_RGB)
        msk_array = np.array(layer_mask)
        for ix in tqdm(range(total_RGB),"Extracting"):
            data_RGB = img_array[ix,...]
            data_msk = msk_array[ix,...]
            
            im_RGB = Image.fromarray(data_RGB)
            im_msk = Image.fromarray(data_msk)
            
            im_RGB1 = im_RGB.save(os.path.join(zip_dir.name,'RGB_'+str(ix)+".png"))
            im_msk1 = im_msk.save(os.path.join(zip_dir.name,'MSK_'+str(ix)+".png"))
        
    elif len(nbr_image)==3:
        assert len(nbr_mask)==2, "NOT A SINGLE MASK"
        img_array = np.array(layer_RGB)
        msk_array = np.array(layer_mask)
        im_RGB = Image.fromarray(img_array)
        im_msk = Image.fromarray(msk_array)
            
        im_RGB1 = im_RGB.save(os.path.join(zip_dir.name,"RGB.png"))
        im_msk1 = im_msk.save(os.path.join(zip_dir.name,"MSK.png"))
    elif len(nbr_image)==2:
        assert len(nbr_mask)>=2, "NOT A SINGLE MASK"
        img_array = np.array(layer_RGB)
        msk_array = np.array(layer_mask)
        im_RGB = Image.fromarray(img_array)
        im_msk = Image.fromarray(msk_array)
            
        im_RGB1 = im_RGB.save(os.path.join(zip_dir.name,"RGB.png"))
        im_msk1 = im_msk.save(os.path.join(zip_dir.name,"MSK.png"))
    else:
        assert len(nbr_image)==4 or len(nbr_mask)==2, "NOT STACK NOR SINGLE IMAGE"
    shutil.make_archive(filename, 'zip', zip_dir.name)
    
    for ix in os.listdir(zip_dir.name):
        os.remove(os.path.join(zip_dir.name,ix))
        
    show_info('Compressed file done')