doggy.machine-learning/classifyDog.py at master · openhack-Doggy/doggy.machine-learning · 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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Apr 26 07:16:04 2018

@author: raghav prabhu
Re-modified TensorFlow classification file according to our need.
"""
import tensorflow as tf
import sys
import os
import csv

# Disable tensorflow compilation warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

'''
Classify images from test folder and predict dog breeds along with score.
'''


def classify_image(image_path, headers):
    f = open('submit.csv', 'w')
    writer = csv.DictWriter(f, fieldnames=headers)
    writer.writeheader()

    # Loads label file, strips off carriage return
    label_lines = [line.rstrip() for line
                   in tf.gfile.GFile("trained_model/retrained_labels.txt")]

    # Unpersists graph from file
    with tf.gfile.FastGFile("trained_model/retrained_graph.pb", 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')

    files = os.listdir(image_path)
    with tf.Session() as sess:
        for file in files:
            # Read the image_data
            image_data = tf.gfile.FastGFile(image_path + '/' + file, 'rb').read()
            # Feed the image_data as input to the graph and get first prediction
            softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')

            predictions = sess.run(softmax_tensor, \
                                   {'DecodeJpeg/contents:0': image_data})

            # Sort to show labels of first prediction in order of confidence
            top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
            records = []
            row_dict = {}
            head, tail = os.path.split(file)
            row_dict['id'] = tail.split('.')[0]

            for node_id in top_k:
                human_string = label_lines[node_id]

                # Some breed names are mismatching with breed name in csv header names.
                human_string = human_string.replace(" ", "_")
                if (human_string == 'german_short_haired_pointer'):
                    human_string = 'german_short-haired_pointer'
                if (human_string == 'shih_tzu'):
                    human_string = 'shih-tzu'
                if (human_string == 'wire_haired_fox_terrier'):
                    human_string = 'wire-haired_fox_terrier'
                if (human_string == 'curly_coated_retriever'):
                    human_string = 'curly-coated_retriever'
                if (human_string == 'black_and_tan_coonhound'):
                    human_string = 'black-and-tan_coonhound'
                if (human_string == 'soft_coated_wheaten_terrier'):
                    human_string = 'soft-coated_wheaten_terrier'
                if (human_string == 'flat_coated_retriever'):
                    human_string = 'flat-coated_retriever'
                score = predictions[0][node_id]
                print('%s (score = %.5f)' % (human_string, score))
                row_dict[human_string] = score
            records.append(row_dict.copy())
            writer.writerows(records)
    f.close()


def main():
    test_data_folder = 'test'

    template_file = open('sample_submission.csv', 'r')
    d_reader = csv.DictReader(template_file)

    # get fieldnames from DictReader object and store in list
    headers = d_reader.fieldnames
    template_file.close()
    classify_image(test_data_folder, headers)


if __name__ == '__main__':
    main()