Skip to content

Commit ffbae5f

Browse files
madimadi
authored
Add files via upload
run_prediction is just the jupyter notebook besides it exports the boxes convert_coords_boxes is a post processing script to convert boxes into spatial coordinates and creates a shapefile
1 parent 061c11f commit ffbae5f

File tree

2 files changed

+298
-0
lines changed

2 files changed

+298
-0
lines changed

‎convert_coords_boxes.py

+154
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
# coding: utf-8
2+
3+
import numpy as np
4+
import os
5+
import sys
6+
import pandas as pd
7+
from osgeo import gdal,ogr,osr
8+
9+
PATH_TO_TEST_IMAGES_DIR = '/home/madi/Projects/models/research/trees_recognition/images/pred/'
10+
TEST_IMAGE_PATHS = os.path.join(PATH_TO_TEST_IMAGES_DIR, 'aaa_pt604000-4399000.jpg' )
11+
TEST_CSV_PATH = os.path.join(PATH_TO_TEST_IMAGES_DIR, 'aaa_pt604000-4399000.csv' )
12+
CSV_REPR_PATH = os.path.join(PATH_TO_TEST_IMAGES_DIR, 'aaa_pt604000-4399000_repr.csv' )
13+
14+
#-------------------------------------------------------------------------------
15+
16+
def GetExtent(gt, cols, rows):
17+
ext = []
18+
xarr = [0, cols]
19+
yarr = [0, rows]
20+
for px in xarr:
21+
for py in yarr:
22+
x = gt[0] + (px * gt[1]) + (py * gt[2])
23+
y = gt[3] + (px * gt[4]) + (py * gt[5])
24+
ext.append([x, y])
25+
print x, y
26+
yarr.reverse()
27+
return ext
28+
29+
#-------------------------------------------------------------------------------
30+
31+
def ReprojectCoords(coords, ext):
32+
'''From normalized coords to src
33+
'''
34+
xmin_img = ext[0][0]
35+
xmax_img = ext[2][0]
36+
ymin_img = ext[1][1]
37+
ymax_img = ext[0][1]
38+
trans_coords=[]
39+
ymax_t = ymin_img + ((ymax_img - ymin_img) * coords[1][1]) #ymax
40+
xmin_t = xmin_img + ((xmax_img - xmin_img) * coords[0][0]) #xmin
41+
ymin_t = ymin_img + ((ymax_img - ymin_img) * coords[0][1]) #ymin
42+
xmax_t = xmin_img + ((xmax_img - xmin_img) * coords[1][0]) #xmax
43+
trans_coords.append([ymax_t, xmin_t, ymin_t, xmax_t])
44+
return trans_coords
45+
46+
#-------------------------------------------------------------------------------
47+
48+
def writeShapefile(df, outShp, src):
49+
driver = ogr.GetDriverByName("ESRI Shapefile")
50+
if os.path.exists(outShp):
51+
driver.DeleteDataSource(outShp)
52+
ds = driver.CreateDataSource(outShp)
53+
layer = ds.CreateLayer('Boxes', src, ogr.wkbPolygon)
54+
# Add fields
55+
field_classe = ogr.FieldDefn("classe", ogr.OFTInteger)
56+
field_classe.SetWidth(3)
57+
layer.CreateField(field_classe)
58+
field_score = ogr.FieldDefn("score", ogr.OFTReal)
59+
field_score.SetWidth(20)
60+
layer.CreateField(field_score)
61+
# Process df and add the attributes and features to the shapefile
62+
for i in range(len(df)):
63+
# Create geometry of a box
64+
box = ogr.Geometry(ogr.wkbLinearRing)
65+
box.AddPoint(df.iloc[i].xmin, df.iloc[i].ymin) #LL
66+
box.AddPoint(df.iloc[i].xmin, df.iloc[i].ymax) #UL
67+
box.AddPoint(df.iloc[i].xmax, df.iloc[i].ymax) #UR
68+
box.AddPoint(df.iloc[i].xmax, df.iloc[i].ymin) #LR
69+
box.AddPoint(df.iloc[i].xmin, df.iloc[i].ymin) #close ring
70+
# Create polygon
71+
poly = ogr.Geometry(ogr.wkbPolygon)
72+
poly.AddGeometry(box)
73+
#print 'Polygon area = ', poly.GetArea()
74+
#print poly.ExportToWkt()
75+
feature = ogr.Feature(layer.GetLayerDefn())
76+
feature.SetField("class", df.iloc[i].classe)
77+
feature.SetField("score", df.iloc[i].score)
78+
# Set the feature geometry using the box
79+
feature.SetGeometry(poly)
80+
# Create the feature in the layer (shapefile)
81+
layer.CreateFeature(feature)
82+
# Flush memory
83+
feature.Destroy()
84+
# Deference all
85+
ds = layer = feature = poly = None
86+
87+
#-------------------------------------------------------------------------------
88+
89+
# create dataframe
90+
df = pd.read_csv(TEST_IMAGE_PATHS.split(".")[0] + ".csv", \
91+
header = 0, \
92+
names = ['ID','ymax', 'xmin', 'ymin', 'xmax', 'classe', 'score'], \
93+
index_col = 'ID', \
94+
usecols = ['ID','ymax', 'xmin', 'ymin', 'xmax', 'classe', 'score'])
95+
96+
# Remove records having score < 0.10
97+
df = df.drop(df[df.score < 0.10].index)
98+
99+
# The y axis is reversed
100+
df['ymin'] = df['ymin'].apply(lambda x: 1 - x)
101+
df['ymax'] = df['ymax'].apply(lambda x: 1 - x)
102+
103+
# Retrieve coordinates and src of image corners
104+
raster = TEST_IMAGE_PATHS
105+
ds = gdal.Open(raster)
106+
gt = ds.GetGeoTransform()
107+
cols = ds.RasterXSize
108+
rows = ds.RasterYSize
109+
ext = GetExtent(gt,cols,rows)
110+
wkt = ds.GetProjection()
111+
src = osr.SpatialReference()
112+
src.ImportFromWkt(wkt)
113+
114+
# write file
115+
fileCSV = open(CSV_REPR_PATH, "w")
116+
fileCSV.write("ID, ymax, xmin, ymin, xmax, class, score" + "\n")
117+
coords = []
118+
reprojected = []
119+
for i in range(len(df)):
120+
coord = [[df.xmin[i],df.ymin[i]], [df.xmax[i],df.ymax[i]]]
121+
coords.append(coord)
122+
# Perform coordinate conversion of boxes
123+
repr = ReprojectCoords(coord, ext)
124+
rep = [i, repr, df.classe[i], df.score[i]]
125+
reprojected.append(rep)
126+
# write csv
127+
fileCSV.write("%s" % i)
128+
fileCSV.write(",")
129+
fileCSV.write("%s" % repr[0][0])
130+
fileCSV.write(",")
131+
fileCSV.write("%s" % repr[0][1])
132+
fileCSV.write(",")
133+
fileCSV.write("%s" % repr[0][2])
134+
fileCSV.write(",")
135+
fileCSV.write("%s" % repr[0][3])
136+
fileCSV.write(",")
137+
fileCSV.write("%s" % df.classe[i])
138+
fileCSV.write(",")
139+
fileCSV.write("%s" % df.score[i])
140+
fileCSV.write("\n")
141+
142+
fileCSV.close()
143+
144+
# Import clean dataframe
145+
# create dataframe
146+
df = pd.read_csv(CSV_REPR_PATH, \
147+
header = 0, \
148+
names = ['ID','ymax', 'xmin', 'ymin', 'xmax', 'classe', 'score'], \
149+
index_col = 'ID', \
150+
usecols = ['ID','ymax', 'xmin', 'ymin', 'xmax', 'classe', 'score'])
151+
152+
# Create shapefile
153+
outShp = os.path.join(PATH_TO_TEST_IMAGES_DIR, TEST_IMAGE_PATHS.split(".")[0]) + ".shp"
154+
writeShapefile(df, outShp, src)

‎run_prediction.py

+144
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# coding: utf-8
2+
3+
'''
4+
This code performs the prediction using TensorFlow object detection API on the
5+
basis of the code presented in the jupyter notebook at
6+
https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb
7+
This code is supposed to run in the folder:
8+
/models/research/object_detection/
9+
'''
10+
11+
# # Imports
12+
13+
import numpy as np
14+
import os
15+
import six.moves.urllib as urllib
16+
import sys
17+
import tarfile
18+
import tensorflow as tf
19+
import zipfile
20+
21+
from collections import defaultdict
22+
from io import StringIO
23+
from PIL import Image
24+
import time
25+
26+
# ## Env setup
27+
# Path to object detection library
28+
sys.path.append("..")
29+
30+
# ## Object detection imports
31+
from utils import label_map_util
32+
from utils import visualization_utils as vis_util
33+
34+
# # Model preparation
35+
36+
# ## Variables
37+
# Any model exported using the `export_inference_graph.py` tool can be loaded
38+
#here simply by changing `PATH_TO_CKPT` to point to a new .pb file.
39+
40+
# model
41+
MODEL_NAME = 'tree_detection_graph'
42+
ROOT_PATH = '/home/madi/Projects/models/research/trees_recognition/training/'
43+
44+
# Path to frozen detection graph. This is the actual model that is
45+
#used for the object detection.
46+
PATH_TO_CKPT = ROOT_PATH + MODEL_NAME + '/frozen_inference_graph.pb'
47+
48+
# List of the strings that is used to add correct label for each box.
49+
PATH_TO_LABELS = os.path.join(ROOT_PATH, 'trees_detection.pbtxt')
50+
51+
NUM_CLASSES = 2
52+
53+
PATH_TO_TEST_IMAGES_DIR = '/home/madi/Projects/models/research/trees_recognition/images/pred/'
54+
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'aaa_pt604000-4399000.jpg' )]
55+
56+
# Size, in inches, of the output images
57+
IMAGE_SIZE = (3470, 3470)
58+
59+
# ## Load a (frozen) Tensorflow model into memory.
60+
detection_graph = tf.Graph()
61+
with detection_graph.as_default():
62+
od_graph_def = tf.GraphDef()
63+
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
64+
serialized_graph = fid.read()
65+
od_graph_def.ParseFromString(serialized_graph)
66+
tf.import_graph_def(od_graph_def, name = '')
67+
68+
69+
# ## Loading label map
70+
# Label maps map indices to category names, so that when our convolution network
71+
#predicts `5`, we know that this corresponds to `airplane`. Here we use internal
72+
#utility functions, but anything that returns a dictionary mapping integers to
73+
#appropriate string labels would be fine
74+
75+
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
76+
categories = label_map_util.convert_label_map_to_categories(label_map, \
77+
max_num_classes = NUM_CLASSES, \
78+
use_display_name = True)
79+
category_index = label_map_util.create_category_index(categories)
80+
81+
# ## Helper code
82+
def load_image_into_numpy_array(image):
83+
(im_width, im_height) = image.size
84+
return np.array(image.getdata()).reshape(
85+
(im_height, im_width, 3)).astype(np.uint8)
86+
87+
# # Detection
88+
start_time = time.time()
89+
with detection_graph.as_default():
90+
with tf.Session(graph = detection_graph) as sess:
91+
# Definite input and output Tensors for detection_graph
92+
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
93+
# Each box represents a part of the image where a particular object was detected.
94+
detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
95+
# Each score represent how level of confidence for each of the objects.
96+
# Score is shown on the result image, together with the class label.
97+
detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
98+
detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
99+
num_detections = detection_graph.get_tensor_by_name('num_detections:0')
100+
for image_path in TEST_IMAGE_PATHS:
101+
image = Image.open(image_path)
102+
# the array based representation of the image will be used later in order
103+
# to prepare the result image with boxes and labels on it.
104+
image_np = load_image_into_numpy_array(image)
105+
# Expand dimensions since the model expects images to have shape:
106+
# [1, None, None, 3]
107+
image_np_expanded = np.expand_dims(image_np, axis = 0)
108+
# Actual detection.
109+
(boxes, scores, classes, num) = sess.run(
110+
[detection_boxes, detection_scores, detection_classes, num_detections],
111+
feed_dict = {image_tensor: image_np_expanded})
112+
113+
# write file
114+
CSVfile = open(image_path.split(".")[0] + ".csv", "w")
115+
# note that y axis is reversed
116+
CSVfile.write("ID, ymax, xmin, ymin, xmax, class, score" + "\n")
117+
for item in range(0, int(num[0])):
118+
CSVfile.write("%s" % item)
119+
CSVfile.write(",")
120+
boxes[0][item].tofile(CSVfile, sep=",", format="%s")
121+
CSVfile.write(",")
122+
classes[0][item].tofile(CSVfile, sep=",", format="%s")
123+
CSVfile.write(",")
124+
scores[0][item].tofile(CSVfile, sep=",", format="%s")
125+
CSVfile.write("\n")
126+
CSVfile.close()
127+
128+
elapsed_time = time.time() - start_time
129+
print "Time for detection and CSV writing ", elapsed_time
130+
131+
# Visualization of the results of a detection.
132+
import matplotlib.pyplot as plt
133+
import matplotlib.image as mpimg
134+
vis_util.visualize_boxes_and_labels_on_image_array(
135+
image_np,
136+
np.squeeze(boxes),
137+
np.squeeze(classes).astype(np.int32),
138+
np.squeeze(scores),
139+
category_index,
140+
use_normalized_coordinates=True,
141+
line_thickness=8)
142+
143+
imgplot = plt.imshow(image_np)
144+
plt.show()

0 commit comments

Comments
 (0)