PAT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import time\n",
    "from PIL import Image,ImageTk\n",
    "import tkinter as TK\n",
    "from tkinter import filedialog\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.backends.backend_tkagg import (\n",
    "    FigureCanvasTkAgg, NavigationToolbar2Tk)\n",
    "from matplotlib.figure import Figure\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "##define the network##\n",
    "\n",
    "class single_conv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(single_conv, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3,stride=1,padding=(1,1))\n",
    "        self.bn = nn.BatchNorm2d(out_ch)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.bn(self.conv(x)))\n",
    "        return x\n",
    "\n",
    "class down(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(down, self).__init__()\n",
    "        self.down = nn.MaxPool2d(2,2)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.down(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "class up(nn.Module): \n",
    "    def __init__(self):\n",
    "        super(up, self).__init__()       \n",
    "        self.up = nn.Upsample(scale_factor=2, mode='bilinear')\n",
    "\n",
    "    def forward(self, x1, x2): \n",
    "        x1 = self.up(x1)\n",
    "        diffY = x2.size()[2] - x1.size()[2]\n",
    "        diffX = x2.size()[3] - x1.size()[3]\n",
    "        x1 = nn.functional.pad(x1, (diffX // 2, diffX - diffX//2,\n",
    "                        diffY // 2, diffY - diffY//2))\n",
    "        x = torch.cat((x2,x1),1)\n",
    "        \n",
    "        return x\n",
    "        \n",
    "class outconv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(outconv, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3,stride=1,padding=(1,1))\n",
    "        self.sig=nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.sig(self.conv(x))\n",
    "        return x\n",
    "class Diceloss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Diceloss,self).__init__()\n",
    "    \n",
    "    def forward(self,predict,label):\n",
    "        loss=torch.sum(predict*label)+1\n",
    "        loss=1-loss/(torch.sum(predict)+torch.sum(label)-loss+2)\n",
    "        return loss\n",
    "\n",
    "class MyNet_4(nn.Module):\n",
    "    def __init__(self,thickness):\n",
    "        super(MyNet_4, self).__init__()\n",
    "        self.conv1=single_conv(3,thickness)\n",
    "        self.conv2=single_conv(thickness,2*thickness)\n",
    "        self.conv3=single_conv(2*thickness,4*thickness)\n",
    "        self.conv4=single_conv(4*thickness,thickness)\n",
    "        self.conv5=single_conv(2*thickness,thickness)\n",
    "        self.conv6=single_conv(4*thickness,4*thickness)\n",
    "        self.conv7=single_conv(8*thickness,2*thickness)\n",
    "        self.down=down()\n",
    "        self.up=up()\n",
    "        self.out=outconv(thickness,1)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "            x1=self.conv1(x)\n",
    "            x2=self.down(x1)\n",
    "            x2=self.conv2(x2)\n",
    "            x3=self.down(x2)\n",
    "            x3=self.conv3(x3)\n",
    "            x4=self.down(x3)\n",
    "            x4=self.conv6(x4)\n",
    "            x3=self.up(x4,x3)\n",
    "            del x4\n",
    "            x3=self.conv7(x3)\n",
    "            x2=self.up(x3,x2)\n",
    "            del x3\n",
    "            x2=self.conv4(x2)\n",
    "            x1=self.up(x2,x1)\n",
    "            del x2\n",
    "            x1=self.conv5(x1)\n",
    "            x1=self.out(x1).squeeze(1)\n",
    "\n",
    "            return x1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "##get the coordinate of the apex from the mask, the output of the network##\n",
    "\n",
    "def get_coordinate(mask,Thre=30,select_mode='mean'):\n",
    "    height,width=mask.shape\n",
    "    x=np.repeat(np.arange(1,width+1)[None,:],height,axis=0)*mask\n",
    "    y=np.repeat(np.arange(1,height+1)[:,None],width,axis=1)*mask\n",
    "    x=x[x!=0]\n",
    "    y=y[y!=0]\n",
    "    mask_new=(np.abs(x-np.median(x))<Thre)*(np.abs(y-np.median(y))<Thre)\n",
    "    x=x[mask_new]\n",
    "    y=y[mask_new]\n",
    "    if select_mode=='median':\n",
    "        coordinate=[np.round(np.median(x)).astype(int)-1,np.round(np.median(y)).astype(int)-1]\n",
    "    else:\n",
    "        coordinate=[np.round(np.mean(x)).astype(int)-1,np.round(np.mean(y)).astype(int)-1]\n",
    "    \n",
    "    return coordinate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "##locate the search range2. ingore the pixel out of the search range2##\n",
    "\n",
    "def locate_search_range(mask_,search_range,coordinate):\n",
    "    search_center=[coordinate[0],coordinate[1]]\n",
    "    search_mask=np.zeros(mask_.shape).astype(int)\n",
    "    search_mask[max(search_center[1]-search_range,0):min(search_center[1]+search_range+1,mask_.shape[0]),\n",
    "                    max(search_center[0]-search_range,0):min(search_center[0]+search_range+1,mask_.shape[1])]=1\n",
    "    mask_=mask_*search_mask\n",
    "    return mask_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "##color filter to remove pixels of non-plant##\n",
    "\n",
    "def color_filter(image):\n",
    "    eb=0.001\n",
    "    mask_g=(image[:,:,1]>50)*((image[:,:,1]/(image[:,:,0]+eb))>1.1)*((image[:,:,1]/(image[:,:,2]+eb))>1.1)\n",
    "    mask_w=((image[:,:,1]>180)*((image[:,:,1]/(image[:,:,0]+eb))>0.9)*((image[:,:,0]/(image[:,:,1]+eb))>0.9)\n",
    "    *((image[:,:,1]/(image[:,:,2]+eb))>0.9)*((image[:,:,2]/(image[:,:,1]+eb))>0.9))\n",
    "    for i in [0,1,2]:\n",
    "        image[:,:,i]=image[:,:,i]*(mask_g+mask_w)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "## determine the search_range1##\n",
    "\n",
    "def deside_search_range():\n",
    "    test_pic=torch.zeros(1,3,401,401).to(device)\n",
    "    t=time.time()\n",
    "    MyModel.forward(test_pic)\n",
    "    t=time.time()-t\n",
    "    if t<0.04:\n",
    "        return 200\n",
    "    elif t>0.17:\n",
    "        return 50\n",
    "    else:\n",
    "        return 100\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "## open the video to analyse and initialize some parameters##\n",
    "\n",
    "def Select_File():\n",
    "    if tracking:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global file_path,label1,selected_video,cap,height,width,image,canvas,displayed_frame,fps_of_video,coordinates,rval,frame\n",
    "    global MyModel,search_range,search_range2,Threshold_pick_piexls,Threshold_outliners,pick_range,confidence_threshold,learnning_rate\n",
    "    global scales,label2,frame_interval,m_standard,update_decay,time_threshold\n",
    "    file_path = filedialog.askopenfilename(title='Select video')\n",
    "    cap = cv2.VideoCapture(file_path)\n",
    "    height,width=int(cap.get(4)),int(cap.get(3))\n",
    "    rval, frame = cap.read()\n",
    "    if rval==True:\n",
    "        if Enable_color_filter.get():\n",
    "            image=ImageTk.PhotoImage(Image.fromarray(color_filter(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))))\n",
    "        else:\n",
    "            image=ImageTk.PhotoImage(Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)))\n",
    "        displayed_frame=canvas.create_image(width//2, height//2, image=image)\n",
    "        selected_video=True\n",
    "        label1.config(text=file_path)\n",
    "        canvas.config(width=width, height=height)\n",
    "        m_standard=torch.load(model_save_path+'MyModel_%d_%d_epoch3_^.pth'%(level,thickness))\n",
    "        MyModel.load_state_dict(m_standard)\n",
    "        fps_of_video=[]\n",
    "        coordinates=[]\n",
    "        search_range=deside_search_range()\n",
    "        search_range2=40\n",
    "        Threshold_pick_piexls=0.75\n",
    "        pick_range=15\n",
    "        Threshold_outliners=pick_range \n",
    "        confidence_threshold=0.95\n",
    "        learnning_rate=0.14\n",
    "        update_decay=0.1\n",
    "        time_threshold=4\n",
    "        scales=None\n",
    "        frame_interval=None\n",
    "        label2.grid_forget()\n",
    "        label3.grid_forget()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "## start to track the apex##\n",
    "\n",
    "def start_track():\n",
    "    global tracking,fps_of_video,coordinates,selection_finished,coordinate,image,rval\n",
    "    global frame,displayed_frame,num,win,lastDraw,stop_tracking,selected_apex,mask,mask_\n",
    "    global m_standard,update_decay,time_threshold\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    tracking=True\n",
    "    button4.grid_forget()\n",
    "    button5.grid(row=1,column=2,sticky=TK.W,padx=20,pady=10)\n",
    "    num=1\n",
    "    while cap.isOpened():\n",
    "        if rval==True:\n",
    "            frame_=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "            if Enable_color_filter.get():\n",
    "                frame_=color_filter(frame_)\n",
    "            image=ImageTk.PhotoImage(Image.fromarray(frame_))\n",
    "            frame_=(torch.tensor(frame_.transpose(2,0,1)[None,:,:,:])/255.0)\n",
    "            if num==1:\n",
    "                frame_=frame_.to(device)\n",
    "            else:\n",
    "                frame_=frame_[:,:,max(coordinates[-1][1]-search_range,0):min(coordinates[-1][1]+search_range+1,height),\n",
    "                                max(coordinates[-1][0]-search_range,0):min(coordinates[-1][0]+search_range+1,width)].to(device)\n",
    "            with torch.no_grad():\n",
    "                mask=MyModel.forward(frame_)[0]\n",
    "                mask_=np.array(mask.cpu()>Threshold_pick_piexls)\n",
    "                if len(coordinates)>0:\n",
    "                    if num==1:\n",
    "                        mask_=locate_search_range(mask_,search_range2,coordinate)\n",
    "                    else:\n",
    "                        mask_=locate_search_range(mask_,search_range2,\n",
    "                                                    [coordinate[0]-max(coordinates[-1][0]-search_range,0),\n",
    "                                                    coordinate[1]-max(coordinates[-1][1]-search_range,0)])\n",
    "            if not selected_apex:\n",
    "                coordinate=get_coordinate(mask_,Threshold_outliners)\n",
    "            else:\n",
    "                mask_=np.zeros((height,width)).astype(int)\n",
    "            if coordinate[0]<0 or coordinate[1]<0 or coordinate[0]>width or coordinate[1]>height :#process if fail to detect the apex\n",
    "                if len(coordinates)>0:\n",
    "                    if num==1:\n",
    "                        coordinate=coordinates[-1]\n",
    "                    else:\n",
    "                        coordinate=[coordinates[-1][0]-max(coordinates[-1][0]-search_range,0),\n",
    "                                    coordinates[-1][1]-max(coordinates[-1][1]-search_range,0)]\n",
    "            else:\n",
    "                block=np.zeros(mask_.shape).astype(int)\n",
    "                block[max(coordinate[1]-pick_range,0):min(coordinate[1]+pick_range+1,height),\n",
    "                      max(coordinate[0]-pick_range,0):min(coordinate[0]+pick_range+1,width)]=1\n",
    "                mask_updater=torch.tensor(block).float().to(device)\n",
    "                Loss=Diceloss()\n",
    "                with torch.no_grad():\n",
    "                    confidence=Loss.forward(mask,mask_updater)\n",
    "                if confidence>confidence_threshold:\n",
    "                    print(confidence)\n",
    "                    Optimizer=torch.optim.SGD(MyModel.parameters(),learnning_rate,momentum=0.9,weight_decay=0.0005)\n",
    "                    t_start=time.time()\n",
    "                    while (1):\n",
    "                        if time.time()-t_start>time_threshold:\n",
    "                            break\n",
    "                        predict=MyModel.forward(frame_)[0]\n",
    "                        loss=Loss(predict,mask_updater)\n",
    "                        confidence=loss.item()\n",
    "                        if  confidence<=confidence_threshold*0.6:\n",
    "                            break\n",
    "                        Optimizer.zero_grad()\n",
    "                        loss.backward()\n",
    "                        Optimizer.step()\n",
    "                        m_now=MyModel.state_dict()\n",
    "                        for key in m_now:\n",
    "                            m_now[key]=(1-update_decay)*m_standard[key].to(device)+update_decay*m_now[key]\n",
    "                        MyModel.load_state_dict(m_now)\n",
    "                    m_standard=m_now\n",
    "            t = time.time()\n",
    "            if num!=1:\n",
    "                coordinate=[coordinate[0]+max(coordinates[-1][0]-search_range,0),coordinate[1]+max(coordinates[-1][1]-search_range,0)]\n",
    "                fps=1/(t-tp)\n",
    "                fps_of_video.append(fps)\n",
    "                canvas.delete(displayed_frame)\n",
    "                canvas.delete(lastDraw)\n",
    "                displayed_frame=canvas.create_image(width//2, height//2, image=image)\n",
    "            if stop_tracking:\n",
    "                tracking=False\n",
    "                stop_tracking=False\n",
    "                break\n",
    "            if selected_apex:\n",
    "                selected_apex=False\n",
    "                canvas.delete(lastDraw)\n",
    "            tp=t\n",
    "            lastDraw = canvas.create_rectangle(max(coordinate[0]-pick_range,0), max(coordinate[1]-pick_range,0)\n",
    "                                    ,min(coordinate[0]+pick_range,width), min(coordinate[1]+pick_range,height), outline='red')\n",
    "            win.update_idletasks()\n",
    "            win.update()\n",
    "            coordinates.append(coordinate)\n",
    "            num+=1\n",
    "        else:\n",
    "            cap.release()\n",
    "            break\n",
    "        rval, frame = cap.read()\n",
    "    tracking=False\n",
    "    button5.grid_forget()\n",
    "    button4.grid(row=1,column=2,sticky=TK.W,padx=20,pady=10)\n",
    "\n",
    "## stop tracking##\n",
    "def stop_track():\n",
    "    global stop_tracking\n",
    "    if tracking:\n",
    "        stop_tracking=True\n",
    "        button5.grid_forget()\n",
    "        button4.grid(row=1,column=2,sticky=TK.W,padx=20,pady=10)\n",
    "    else:\n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions is used to draw a box in the canvas to select the apex ##\n",
    "\n",
    "\n",
    "def Select_apex1():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection\n",
    "    enable_selection=True\n",
    "    button2.grid_forget()\n",
    "    button3.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    \n",
    "def Select_apex2():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection,coordinate,selected_apex,selection_finished,pick_range\n",
    "    enable_selection=False\n",
    "    button3.grid_forget()\n",
    "    button2.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    if selection_finished:\n",
    "        coordinate=[(selected_area[0]+selected_area[1])//2,(selected_area[2]+selected_area[3])//2]\n",
    "        if selected_area[1]>selected_area[0] and selected_area[3]>selected_area[2]:\n",
    "            pick_range=(selected_area[1]-selected_area[0]+selected_area[3]-selected_area[2])//4\n",
    "            Threshold_outliners=pick_range\n",
    "        selected_apex=True\n",
    "        selection_finished=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions is used to draw a box in the canvas to decide the search range 1 ##\n",
    "\n",
    "def search_range_1_1():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection2\n",
    "    enable_selection2=True\n",
    "    button6.grid_forget()\n",
    "    button7.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    \n",
    "def search_range_1_2():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection2,selection_finished,search_range\n",
    "    enable_selection2=False\n",
    "    button7.grid_forget()\n",
    "    button6.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    if selection_finished:\n",
    "        if selected_area[1]>selected_area[0] and selected_area[3]>selected_area[2]:\n",
    "            search_range=(selected_area[1]-selected_area[0]+selected_area[3]-selected_area[2])//4\n",
    "        canvas.delete(lastDraw)\n",
    "        selection_finished=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions is used to draw a box in the canvas to decide the search range 2 ##\n",
    "\n",
    "def search_range_2_1():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection3\n",
    "    enable_selection3=True\n",
    "    button8.grid_forget()\n",
    "    button9.grid(row=2,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    \n",
    "def search_range_2_2():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection4:\n",
    "        return\n",
    "    global enable_selection3,selection_finished,search_range2\n",
    "    enable_selection3=False\n",
    "    button9.grid_forget()\n",
    "    button8.grid(row=2,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    if selection_finished:\n",
    "        if selected_area[1]>selected_area[0] and selected_area[3]>selected_area[2]:\n",
    "            search_range2=(selected_area[1]-selected_area[0]+selected_area[3]-selected_area[2])//4\n",
    "        canvas.delete(lastDraw)\n",
    "        selection_finished=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions is used to map the scales in frames to in the real world ##\n",
    "\n",
    "def Scale_1():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3:\n",
    "        return\n",
    "    global enable_selection4\n",
    "    enable_selection4=True\n",
    "    button10.grid_forget()\n",
    "    button11.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    entry1.delete(0,\"end\")\n",
    "    entry1.grid(row=2,column=1,sticky=TK.W,padx=20)\n",
    "    label2.grid_forget()   \n",
    "def Scale_2():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3:\n",
    "        return\n",
    "    global enable_selection4,selection_finished,scales\n",
    "    enable_selection4=False\n",
    "    button11.grid_forget()\n",
    "    entry1.grid_forget()\n",
    "    button10.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    if selection_finished:\n",
    "        try:\n",
    "            scales=float(entry1.get())/np.sqrt((selected_area[1]-selected_area[0])**2+(selected_area[3]-selected_area[2])**2)           \n",
    "        except:\n",
    "            pass \n",
    "        canvas.delete(lastDraw)\n",
    "        selection_finished=False\n",
    "    if scales!=None:\n",
    "        label2.config(text=\"Scales = %.3f mm/pixel\"%(scales))\n",
    "        label2.grid(row=2,column=1,sticky=TK.W,padx=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions is used to map the frame interval to the time in the real world ##\n",
    "\n",
    "def Frame_interval_1():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    button12.grid_forget()\n",
    "    button13.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    entry2.delete(0,\"end\")\n",
    "    entry2.grid(row=2,column=1,sticky=TK.W,padx=20)\n",
    "    label3.grid_forget()   \n",
    "def Frame_interval_2():\n",
    "    if tracking or not selected_video:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global frame_interval\n",
    "    button13.grid_forget()\n",
    "    entry2.grid_forget()\n",
    "    button12.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "    try:\n",
    "        frame_interval=float(entry2.get())         \n",
    "    except:\n",
    "        pass \n",
    "    if frame_interval!=None:\n",
    "        label3.config(text=\"Frame interval = %.1f s/frame\"%(frame_interval))\n",
    "        label3.grid(row=2,column=1,sticky=TK.W,padx=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "## These functions help some functions above draw the figure in the canvas ##\n",
    "\n",
    "def onLeftButtonDown(event):\n",
    "    if tracking:\n",
    "        return\n",
    "    elif not enable_selection and not enable_selection2 and not enable_selection3 and not enable_selection4:\n",
    "        return\n",
    "    global s_X,s_Y,selecting\n",
    "    s_X = TK.IntVar(value=0)\n",
    "    s_Y = TK.IntVar(value=0)\n",
    "    s_X.set(event.x)\n",
    "    s_Y.set(event.y)\n",
    "    selecting = True\n",
    "def onLeftButtonMove(event):\n",
    "    global lastDraw\n",
    "    if (not selecting) or tracking:\n",
    "        return\n",
    "    if enable_selection:\n",
    "        color='red'\n",
    "    elif enable_selection2:\n",
    "        color='yellow'\n",
    "    elif enable_selection3:\n",
    "        color='green'\n",
    "    elif enable_selection4:\n",
    "        color='red'\n",
    "    else:\n",
    "        return\n",
    "    try:\n",
    "        canvas.delete(lastDraw)\n",
    "    except:\n",
    "        pass\n",
    "    if enable_selection4:\n",
    "        lastDraw = canvas.create_line(s_X.get(), s_Y.get(), event.x, event.y, fill=color)\n",
    "    else:\n",
    "        lastDraw = canvas.create_rectangle(s_X.get(), s_Y.get(), event.x, event.y, outline=color)\n",
    "def onLeftButtonUp(event):\n",
    "    global selecting,selected_area,enable_selection,selection_finished\n",
    "    if tracking:\n",
    "        return\n",
    "    elif not enable_selection and not enable_selection2 and not enable_selection3 and not enable_selection4:\n",
    "        return\n",
    "    selecting = False\n",
    "    selection_finished= True\n",
    "    myleft, myright = sorted([s_X.get(), event.x])\n",
    "    mytop, mybottom = sorted([s_Y.get(), event.y])\n",
    "    selected_area=(myleft,myright,mytop,mybottom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "##plot the graphs of the result ##\n",
    "\n",
    "def create_fig(fig):\n",
    "    global graph_canvas,graph_toolbar\n",
    "    graph_canvas = FigureCanvasTkAgg(fig, master=graph)\n",
    "    graph_toolbar = NavigationToolbar2Tk(graph_canvas, graph)\n",
    "    graph_toolbar.update()\n",
    "    graph_canvas.get_tk_widget().pack(side=TK.TOP, fill=TK.BOTH, expand=1)\n",
    "    graph_toolbar.pack(side=TK.BOTTOM, fill=TK.X)\n",
    "    return\n",
    "    \n",
    "def X_vs_T():\n",
    "    global graph_canvas,graph_toolbar\n",
    "    try:\n",
    "        graph_canvas.get_tk_widget().destroy()\n",
    "        graph_toolbar.destroy()\n",
    "    except:\n",
    "        pass\n",
    "    create_fig(fig1)\n",
    "    return\n",
    "\n",
    "def Y_vs_T():\n",
    "    global graph_canvas,graph_toolbar\n",
    "    try:\n",
    "        graph_canvas.get_tk_widget().destroy()\n",
    "        graph_toolbar.destroy()\n",
    "    except:\n",
    "        pass\n",
    "    create_fig(fig2)\n",
    "    return\n",
    "\n",
    "def Y_vs_X():\n",
    "    global graph_canvas,graph_toolbar\n",
    "    try:\n",
    "        graph_canvas.get_tk_widget().destroy()\n",
    "        graph_toolbar.destroy()\n",
    "    except:\n",
    "        pass\n",
    "    create_fig(fig3)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "##plot the graphs of the result ##\n",
    "\n",
    "def Plot_the_graphs():\n",
    "    if tracking or not selected_video or len(coordinates)==0:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global graph_canvas,graph_toolbar,graph,fig1,fig2,fig3\n",
    "    positions=np.array(coordinates)\n",
    "    positions[:,1]=height-positions[:,1]\n",
    "    Times=np.arange(0, len(coordinates))\n",
    "    unit_s=\"pixel\"\n",
    "    unit_t=\"frame\"\n",
    "    if scales!=None:\n",
    "        positions=positions*scales\n",
    "        unit_s=\"mm\"\n",
    "    if frame_interval!=None:\n",
    "        Times=Times*frame_interval\n",
    "        unit_t=\"second\"\n",
    "    fig1 = Figure(figsize=(8, 4), dpi=100)\n",
    "    fig1.add_subplot(111,xlabel=\"Time (%s)\"%unit_t,ylabel=\"X (%s)\"%unit_s).plot(Times, positions[:,0])\n",
    "    fig2 = Figure(figsize=(8, 4), dpi=100)\n",
    "    fig2.add_subplot(111,xlabel=\"Time (%s)\"%unit_t,ylabel=\"Y (%s)\"%unit_s).plot(Times, positions[:,1])\n",
    "    fig3 = Figure(figsize=(6, 6), dpi=100)\n",
    "    fig3.add_subplot(111,xlabel=\"X (%s)\"%unit_s,ylabel=\"Y (%s)\"%unit_s).plot(positions[:,0], positions[:,1],'o')\n",
    "    graph = TK.Toplevel(win)\n",
    "    graph.title(\"The movement graph of the apex\")\n",
    "    graph_boxframe1 = TK.Frame(graph)\n",
    "    graph_button1 = TK.Button(graph_boxframe1 ,text=\"X vs T\",command=X_vs_T)\n",
    "    graph_button2 = TK.Button(graph_boxframe1 ,text=\"Y vs T\",command=Y_vs_T)\n",
    "    graph_button3 = TK.Button(graph_boxframe1 ,text=\"Y vs X\",command=Y_vs_X)\n",
    "    graph_button1.grid(row=1,column=1)\n",
    "    graph_button2.grid(row=1,column=2)\n",
    "    graph_button3.grid(row=1,column=3)\n",
    "    graph_boxframe1.pack(side=TK.TOP, fill=TK.X)\n",
    "    graph.grab_set()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "##save the result of tracking##\n",
    "\n",
    "def Save_the_result():\n",
    "    if tracking or not selected_video or len(coordinates)==0:\n",
    "        return\n",
    "    elif enable_selection or enable_selection2 or enable_selection3 or enable_selection4:\n",
    "        return\n",
    "    global graph_canvas,graph_toolbar,graph\n",
    "    positions=np.array(coordinates)\n",
    "    positions[:,1]=height-positions[:,1]\n",
    "    Times=np.arange(0, len(coordinates))\n",
    "    unit_s=\"pixel\"\n",
    "    unit_t=\"frame\"\n",
    "    if scales!=None:\n",
    "        positions=positions*scales\n",
    "        unit_s=\"mm\"\n",
    "    if frame_interval!=None:\n",
    "        Times=Times*frame_interval\n",
    "        unit_t=\"second\"\n",
    "    data=np.hstack((Times[:,None],positions))\n",
    "    csv_file=pd.DataFrame(data,columns=[\"Time (%s)\"%unit_t,\"X (%s)\"%unit_s,\"Y (%s)\"%unit_s])\n",
    "    save_path=filedialog.asksaveasfilename(title='Save the result',filetypes=[('CSV', '*.csv')])\n",
    "    csv_file.to_csv(save_path, sep=\",\",index=False)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#create the network and load parameters#\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model_save_path='./model/'\n",
    "thickness=32\n",
    "level=4\n",
    "try:\n",
    "    MyModel=MyNet_4(32).to(device)\n",
    "except:\n",
    "    device = torch.device(\"cpu\")\n",
    "    MyModel=MyNet_4(32).to(device)\n",
    "MyModel.load_state_dict(torch.load(model_save_path+'MyModel_%d_%d_epoch3_^.pth'%(level,thickness))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create the UI#\n",
    "\n",
    "win=TK.Tk()        \n",
    "win.title('Plant Apex Track')\n",
    "fps_of_video=[]\n",
    "coordinates=[]\n",
    "search_range=100\n",
    "search_range2=40\n",
    "Threshold_pick_piexls=0.75\n",
    "pick_range=15\n",
    "Threshold_outliners=pick_range \n",
    "confidence_threshold=0.95\n",
    "learnning_rate=0.14\n",
    "update_decay=0.1\n",
    "time_threshold=4\n",
    "scales=None\n",
    "frame_interval=None\n",
    "selected_video=False\n",
    "enable_selection=False\n",
    "enable_selection2=False\n",
    "enable_selection3=False\n",
    "enable_selection4=False\n",
    "selected_apex=False\n",
    "selection_finished=False\n",
    "selecting=False\n",
    "tracking=False\n",
    "stop_tracking=False\n",
    "Enable_color_filter=TK.IntVar()\n",
    "file_path=\"Please select a video to track\"\n",
    "boxframe1 = TK.Frame(win, relief=\"sunken\")\n",
    "boxframe2 = TK.Frame(win, relief=\"sunken\",borderwidth=1)\n",
    "boxframe3 = TK.Frame(win, relief=\"sunken\",borderwidth=1)\n",
    "boxframe4 = TK.Frame(win, relief=\"sunken\",borderwidth=1)\n",
    "boxframe5 = TK.Frame(boxframe3)\n",
    "boxframe6 = TK.Frame(boxframe3)\n",
    "label1 = TK.Label(boxframe1, text = file_path)\n",
    "label2= TK.Label(boxframe5)\n",
    "label3= TK.Label(boxframe6)\n",
    "button1 = TK.Button(boxframe1 ,text=\"Select video\",command=Select_File)\n",
    "button2 = TK.Button(boxframe4 ,text=\"Select apex\",command=Select_apex1)\n",
    "button3 = TK.Button(boxframe4 ,text=\"Confirm\",command=Select_apex2)\n",
    "button4 = TK.Button(boxframe4 ,text=\"Track\",command=start_track)\n",
    "button5=TK.Button(boxframe4 ,text=\"Stop\",command=stop_track)\n",
    "button6=TK.Button(boxframe3 ,text=\"Search range 1\",command=search_range_1_1)\n",
    "button7=TK.Button(boxframe3 ,text=\"Confirm\",command=search_range_1_2)\n",
    "button8=TK.Button(boxframe3 ,text=\"Search range 2\",command=search_range_2_1)\n",
    "button9=TK.Button(boxframe3 ,text=\"Confirm\",command=search_range_2_2)\n",
    "button10=TK.Button(boxframe5 ,text=\"Scale(draw a line)\",command=Scale_1)\n",
    "button11=TK.Button(boxframe5 ,text=\"Enter the length (mm)\",command=Scale_2)\n",
    "button12=TK.Button(boxframe6 ,text=\"Frame interval\",command=Frame_interval_1)\n",
    "button13=TK.Button(boxframe6 ,text=\"Enter the interval (second)\",command=Frame_interval_2)\n",
    "button14=TK.Button(boxframe4 ,text=\"Plot the graphs\",command=Plot_the_graphs)\n",
    "button15=TK.Button(boxframe4 ,text=\"Save the result\",command=Save_the_result)\n",
    "Checkbutton1 = TK.Checkbutton(boxframe3, text='Enable color filter', variable=Enable_color_filter, onvalue=1, offvalue=0,)\n",
    "entry1=TK.Entry(boxframe5,width=14)\n",
    "entry2=TK.Entry(boxframe6,width=14)\n",
    "canvas = TK.Canvas(boxframe2)\n",
    "canvas.bind('<Button-1>', onLeftButtonDown)\n",
    "canvas.bind('<B1-Motion>', onLeftButtonMove)\n",
    "canvas.bind('<ButtonRelease-1>', onLeftButtonUp)\n",
    "win.grid_columnconfigure(0,weight=0)\n",
    "win.grid_rowconfigure(0,weight=0)\n",
    "canvas.grid(row=1,column=1,padx=20,pady=10)\n",
    "label1.grid(row=1,column=2,sticky=TK.W,padx=20,pady=10)\n",
    "button1.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button2.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button4.grid(row=1,column=2,sticky=TK.W,padx=20,pady=10)\n",
    "button6.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button8.grid(row=2,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button10.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button12.grid(row=1,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "button14.grid(row=1,column=3,sticky=TK.W,padx=20,pady=10)\n",
    "button15.grid(row=1,column=4,sticky=TK.W,padx=20,pady=10)\n",
    "Checkbutton1.grid(row=5,column=1,sticky=TK.W,padx=20,pady=10)\n",
    "boxframe1.grid(row=1,column=1,columnspan=2,sticky=TK.EW)\n",
    "boxframe2.grid(row=2,column=1,sticky=TK.W)\n",
    "boxframe3.grid(row=2,column=2,sticky=TK.NS)\n",
    "boxframe4.grid(row=3,column=1,columnspan=2,sticky=TK.EW)\n",
    "boxframe5.grid(row=3,column=1,sticky=TK.EW)\n",
    "boxframe6.grid(row=4,column=1,sticky=TK.EW)\n",
    "win.mainloop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}