ct/create_lists.py

import random
import os

exs = []
for f in os.listdir("imgs"):
    ex = f.split(".jpg")[0]
    exs.append(ex)

random.shuffle(exs)

fp = open("train.csv")

train_exs = []

for line in fp:
    train_exs.append(line.strip().split(",")[0])
fp.close()

test_exs = [ex for ex in exs if ex not in train_exs]
print(len(exs), len(test_exs))
#train_exs = exs[:3896]
#test_exs = exs[int(len(exs)*0.2) + 1:]

#train_fp = open("train.csv", "w")
#for ex in train_exs:
#    train_fp.write(ex + ",1\n")
#train_fp.close()

test_fp = open("test.csv", "w")
for ex in test_exs:
    test_fp.write(ex + ",1\n")
test_fp.close()