Segmentation.py 4.36 KB
Newer Older
luroth's avatar
luroth committed
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
'''
Created on 16.03.2017

@author: luroth
'''

#from pylab import plot,show
from numpy import vstack,array
from numpy.random import rand
import numpy.ma as ma
import numpy as np
from scipy.cluster.vq import kmeans,vq
from scipy import misc
#import sklearn.cluster
import matplotlib.pyplot as plt
from osgeo import gdalnumeric
import os
import sys
import getopt
import csv
from scipy import ndimage
from Common import create_GeoTiff
import gdal
import colorsys

def main(argv):
    camera_position_file = ''
    metadata_folder = ''
    output_directory = ''
    
    
    try:
        opts, args = getopt.getopt(argv,"hc:m:p:s:o:")
    except getopt.GetoptError:
        print('test.py -c <camera position file> -m <metadata folder> -o <Output directory>')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('test.py -c <camera position file> -m <metadata folder> -o <Output directory>')
            sys.exit()
        elif opt == "-c":
            camera_position_file = arg
        elif opt == "-m":
            metadata_folder = arg    
        elif opt == "-o":
            output_directory = arg                   
    
    print('Metdadata folder is', metadata_folder)


    with open(camera_position_file, 'r') as csvfile:
        spamreader = csv.reader(csvfile, delimiter='\t', quotechar='|')
        next(spamreader, None)
        next(spamreader, None)
 
        for row in spamreader:  
            PhotoID = os.path.splitext(row[0])[0]
            
            image = misc.imread(os.path.join(metadata_folder, '../../RGB', PhotoID + '.JPG'))
            
            if os.path.isfile(os.path.join(metadata_folder, PhotoID + '_mask_lot.tif')):
                image_mask_all =  misc.imread(os.path.join(metadata_folder, PhotoID + '_mask_lot.tif'))
                image_mask_soil =  misc.imread(os.path.join(metadata_folder, PhotoID + '_mask_soil.tif'))
            
                origin_shape = image.shape
                
                image = image.reshape(-1,3)
                image_mask_all = image_mask_all.reshape(-1)
                image_mask_soil = image_mask_soil.reshape(-1)
                
                image[image_mask_all==0,:] = 0
                
                # Colors
                R = image[:,0]/255.0
                G = image[:,1]/255.0
                B = image[:,2]/255.0

                # ExG
                r= np.divide(R, (np.mean(R)+np.mean(G)+np.mean(B)), dtype=float)
                g= np.divide(G, (np.mean(R)+np.mean(G)+np.mean(B)), dtype=float)
                b= np.divide(B, (np.mean(R)+np.mean(G)+np.mean(B)), dtype=float)
                
                ExG = 2.0* g - r - b
                ExR = 1.4 * r - b
                
                misc.imsave(os.path.join(output_directory, PhotoID + '_ExG.tif'), ExG.reshape(origin_shape[0], origin_shape[1]))
                misc.imsave(os.path.join(output_directory, PhotoID + '_ExR.tif'), ExR.reshape(origin_shape[0], origin_shape[1]))
                
                ExG_ExR = ExG - ExR
                
                plants = np.array(ExG_ExR, dtype=float)
                plants = plants.reshape(origin_shape[0], origin_shape[1])
                
                #gdalnumeric.SaveArray(1*(plants.reshape(origin_shape[0], origin_shape[1])), os.path.join(output_directory, PhotoID + '_segmented_ExG_ExR.tif'), format='GTiff')
                create_GeoTiff(os.path.join(output_directory, PhotoID + '_segmented_ExG_ExR'), plants, gdal.GDT_Float32)
                
                #ExG_soil = ExG[image_mask_soil==1]
                #print(ExG_soil.size)
                #if ExG_soil.size ==0:
                #    continue
                
                #ExG_threshold = np.percentile(ExG_soil, 99.95)
                
                #soil = np.array(ExG > ExG_threshold)
                
                #gdalnumeric.SaveArray(1*(soil.reshape(origin_shape[0], origin_shape[1])), os.path.join(output_directory, PhotoID + '_segmented_ExG.tif'), format='GTiff')
                #misc.imsave(os.path.join(output_directory, PhotoID + '_segmented_ExG.tif'), soil.reshape(origin_shape[0], origin_shape[1]))
                
                
                print(PhotoID, "done")

if __name__ == "__main__":
    main(sys.argv[1:])
    
    print("\n============\nAll images segmented, END\n")