forked from nitini/eas_499_code
/
gen_submission_file.py
62 lines (57 loc) · 1.82 KB
/
gen_submission_file.py
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
import matplotlib
matplotlib.use('Agg')
import caffe
import lmdb
import sys
import time
import csv
import subprocess
import numpy as np
import os
def setup_submission_file(train_folder, file_name):
header = ['image']
p_classes = os.listdir(train_folder)
p_classes.sort()
for p_class in p_classes:
header.append(p_class)
submission_file = open(file_name, 'wb')
submission_writer = csv.writer(submission_file)
submission_writer.writerow(header)
submission_file.close()
def main():
MODEL_FILE = sys.argv[1]
PRETRAINED = sys.argv[2]
mean_file = sys.argv[3]
lmdb_folder = sys.argv[4]
train_folder = sys.argv[5]
seaNet = caffe.Net(MODEL_FILE, PRETRAINED, caffe.TEST)
caffe.set_mode_gpu()
image_mean = np.load(mean_file)
file_name = 'seaNet_submission_' + ('%0.f' % time.time()) + '.csv'
setup_submission_file(train_folder, file_name)
submission_file = open(file_name, 'a')
submission_writer = csv.writer(submission_file)
env = lmdb.open(lmdb_folder)
txn = env.begin()
cursor = txn.cursor()
count = 0
for key, value in cursor:
count += 1
if count % 500 == 0:
print 'Number of Images Processed: ' + str(count)
datum = caffe.proto.caffe_pb2.Datum()
datum.ParseFromString(value)
label = datum.label
image = caffe.io.datum_to_array(datum)
image = image.astype(np.uint8)
image = image - image_mean
image = image * 0.00390625
result = seaNet.forward_all(data=np.array([image]))
probs = result['prob'][0]
img_row = [ '_'.join(key.split('_')[1:])]
img_row.extend(probs)
submission_writer.writerow(img_row)
submission_file.close()
#subprocess.call('mv ' + file_name + ' ../../data_files/', shell=True)
if __name__ == "__main__":
main()