mtracker.github.io/track_video.py at master · lingjzhu/mtracker.github.io · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun May  6 20:29:54 2018

@author: luke
"""
import sys
sys.path.insert(0, './src')

import argparse
import skimage.io as io
import numpy as np
import cnn_model as cnn
import cnn_predict as cp
import pandas as pd
import time
import imageio
import csv
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
import DenseNet
from skimage.morphology import skeletonize
from skimage.transform import resize


parser = argparse.ArgumentParser()

parser.add_argument("-v","--path_to_video",dest='v', help="path to the folder containing the demo video",
                    default="./demo/demo_video.mp4")
parser.add_argument("-t","--model_type",dest="t",help="specify which type of model to use (u for unet and du for dense unet)",default="du")
parser.add_argument("-m","--model",dest="m",help="path to the model",default="./model/dense_compound.hdf5")
parser.add_argument("-o","--output_folder",dest="o",help="path to the output folder",default="./demo/output.csv")
parser.add_argument("-f","--figure",dest="f",help="path to save figures",default="./demo/")
parser.add_argument("-b","--boundary",dest="b",nargs=4, help="the boundary for cropping the central portion of images",default=[106,385,121,487])
parser.add_argument("-n","--n_frames",dest="n",help="plot every Nth frame (default 100)",default=100)


args = parser.parse_args()


# please specify the path to input and output files
path_to_video = args.v

path_to_model = args.m

path_to_csv_output = args.o

path_to_figures = args.f

# initiate the model

if args.t == "du":
    model = DenseNet.DenseUnet_v2(weights=None,input_shape=[128,128,3],loss = "compound")
    model.load_weights(path_to_model)
elif args.t == "u":
    model = cnn.Unet()
    model.initiate(128,128,3)
    model.load(path_to_model)
else:
    print("specify which type of model to use (u for unet and du for dense unet)")

# load video
video = imageio.get_reader(path_to_video)

# initiate an empty dataframe
cnn_prediction = pd.DataFrame(columns=['x','y','uniqueframe'])

# plot every Nth frame
N = int(args.n)

# set the boundary for cropping
boundary = args.b
output_size = [boundary[1]-boundary[0],boundary[3]-boundary[2]]
# get splines

with open(path_to_csv_output, mode='w') as out:
    out = csv.writer(out, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
    out.writerow(['x','y','uniqueframe'])
    for index, image in tqdm(enumerate(video)):
        try:
            img = cp.crop(image,boundary)
            img = cp.normalize(img,[128,128])
            img = np.expand_dims(img,0)
            #img = np.expand_dims(img,-1)
            pred = model.predict(img)
            pred = np.squeeze(pred)
            img = np.squeeze(img)
            pred = resize(np.squeeze(pred),output_size)

            sample = np.where(pred < 0.5,0,1)
            skeleton = skeletonize(sample)
            s = cp.interpolate_tongue_spline(skeleton,smooth=True)

            # convert the coordinate back to fit the original image
            s[:,0]=s[:,0]+boundary[2]
            s[:,1]=s[:,1]+boundary[0]

            #ac = cp.get_active_contour(s,img,smooth=False)
            for j in range(len(s)):
                out.writerow([s[j,0],s[j,1],index])

            # plot every Nth frame for inspection
            if index%N == 0:
                fig, ax = plt.subplots(figsize=(9, 5))
                ax.imshow(image, cmap=plt.cm.gray)
                #ax.plot(three['xcoord'], three['ycoord'], '--r', lw=3)
                ax.plot(s[:,0], s[:,1], 'ro', lw=3)
                fig.savefig(path_to_figures+str(index)+'.jpg')
                plt.clf()
                plt.close()
        except:
            continue