Revision 8dea60e7
Added by dsorber about 13 years ago
| similar images/build_index_faster.py | ||
|---|---|---|
|
#!/opt/local/bin/python2.7
|
||
|
import bz2
|
||
|
import cPickle
|
||
|
import imghdr
|
||
|
import math
|
||
|
import multiprocessing
|
||
|
import os
|
||
|
import Queue
|
||
|
import random
|
||
|
import sys
|
||
|
|
||
|
import pHash
|
||
|
|
||
|
from image import Image
|
||
|
|
||
|
NUM_WORKERS = 7
|
||
|
|
||
|
# Build index
|
||
|
# TODO: write more schtuff here
|
||
|
|
||
|
class IndexWorker(multiprocessing.Process):
|
||
|
|
||
|
def __init__(self, file_list, queue):
|
||
|
multiprocessing.Process.__init__(self)
|
||
|
self.file_list = file_list
|
||
|
self.queue = queue
|
||
|
|
||
|
def run(self):
|
||
|
|
||
|
list_size = len(self.file_list)
|
||
|
ctr = 0
|
||
|
for file_path in self.file_list:
|
||
|
ctr += 1
|
||
|
percent_done = int((float(ctr) / list_size) * 100)
|
||
|
# We can only index jpegs right now
|
||
|
if not imghdr.what(file_path) == 'jpeg':
|
||
|
continue
|
||
|
print '{:d} -- {:3d}% -- {:s}'.format(self.pid, percent_done,
|
||
|
file_path)
|
||
|
image_obj = Image(file_path)
|
||
|
|
||
|
self.queue.put(image_obj, True)
|
||
|
|
||
|
|
||
|
|
||
|
def main():
|
||
|
|
||
|
# Make sure at least one command line argument was provided
|
||
|
if len(sys.argv) < 2:
|
||
|
print 'Usage: build_index_faster.py <path>'
|
||
|
return -1;
|
||
|
|
||
|
# Print out the path we're walking
|
||
|
path = sys.argv[1]
|
||
|
print 'Building image index from path: {:s}\n'.format(path)
|
||
|
|
||
|
master_file_list = []
|
||
|
image_index = []
|
||
|
|
||
|
# Walk the root path and create a list of all files in the root path
|
||
|
for root, dirs, files in os.walk(path):
|
||
|
for file in files:
|
||
|
master_file_list.append(os.path.join(root, file))
|
||
|
|
||
|
random.shuffle(master_file_list)
|
||
|
|
||
|
print 'Done counting, {:d} files total.'.format(len(master_file_list))
|
||
|
|
||
|
workers = []
|
||
|
queue = multiprocessing.Queue(64)
|
||
|
|
||
|
def chunker(seq, size):
|
||
|
return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))
|
||
|
|
||
|
chunk_size = int(math.ceil(len(master_file_list) / NUM_WORKERS))
|
||
|
|
||
|
# Create worker processes
|
||
|
for work_chunk in chunker(master_file_list, chunk_size):
|
||
|
workers.append(IndexWorker(work_chunk, queue))
|
||
|
|
||
|
# Start workers
|
||
|
for worker in workers:
|
||
|
worker.start()
|
||
|
|
||
|
print 'Waiting for workers to complete'
|
||
|
|
||
|
# Loop until all workers are finished
|
||
|
while any([w.is_alive() for w in workers]):
|
||
|
try:
|
||
|
image_obj = queue.get(True, 2)
|
||
|
image_index.append(image_obj)
|
||
|
except Queue.Empty:
|
||
|
print 'Queue was empty'
|
||
|
continue
|
||
|
|
||
|
print 'Workers are complete: {:d} images found'.format(len(image_index))
|
||
|
|
||
|
# Write index as a bzip2'd file to save space
|
||
|
index_file = bz2.BZ2File('images_index.bz2', 'w')
|
||
|
cPickle.dump(image_index, index_file)
|
||
|
index_file.close()
|
||
|
|
||
|
print '\nIndex file images_index.bz2 created. Done.'
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
sys.exit(main())
|
||
I wrote a multi process version of the image indexer. It is MUCH faster but I keep running into problematic jpegs that cause pHash and CImg to crash. Unfortunately it doesn't seem like there is a graceful way to handle the crash. I will play with this more as I have time.