benob · November 6, 2022 12:18
diff --git a/gallery.py b/gallery.py
 # Searchable image gallery using CLIP vectors to represent picture content and text queries
 # * Install requirements:
 # pip install annoy pyvips git+https://github.com/Lednik7/CLIP-ONNX.git git+https://github.com/openai/CLIP.git bottle protobuf==3.20
 # * Download models:
 # wget https://clip-as-service.s3.us-east-2.amazonaws.com/models/onnx/ViT-B-32/{visual,textual}.onnx
 # * Run server and open http://127.0.0.1:8080:
 # python gallery.py /path/to/pictures/directory
 import os
 import sys
 import glob
 import struct
 import hashlib
 import sqlite3

 import clip
 import pyvips
 import numpy as np
 from PIL import Image
 from annoy import AnnoyIndex
 from clip_onnx import clip_onnx
 from bottle import run, get, static_file, redirect, request

 directory = sys.argv[1]

 os.makedirs('thumbnails', exist_ok=True)
 os.makedirs('features', exist_ok=True)

 # paths are stored in a simple sqlite db
 con = sqlite3.connect('db.sqlite')
 cur = con.cursor()
 cur.execute('CREATE TABLE IF NOT EXISTS pictures (path TEXT, date DATE, width INT, height INT, thumbnail TEXT, features TEXT)')
 con.commit()

 # load clip models
 _, preprocess = clip.load("ViT-B/32", device="cpu", jit=False)
 onnx_model = clip_onnx(None)
 onnx_model.load_onnx(visual_path="visual.onnx", textual_path="textual.onnx", logit_scale=100.0) # model.logit_scale.exp()
 onnx_model.start_sessions(providers=["CPUExecutionProvider"])

 # generate a vector of features with clip
 def extract_features(path):
  image_input = Image.open(path)
  image = preprocess(image_input).unsqueeze(0).cpu() # [1, 3, 224, 224]
  image_onnx = image.detach().cpu().numpy().astype(np.float32)
  image_features = onnx_model.encode_image(image_onnx)
  return image_features[0]

 # process a single picture
 def process(path):
  # pass already processed files
  if (1,) == cur.execute('SELECT 1 FROM pictures WHERE path = ?', (path,)).fetchone():
    return

  # compute unique identifier from path
  name = hashlib.sha1(path.encode('utf8')).hexdigest()

  # generate and save thumbnail
  image = pyvips.Image.new_from_file(path, access='sequential')
  image = image.autorot()
  date = image.get("exif-ifd0-DateTime")
  thumb = pyvips.Image.thumbnail_image(image, 224, crop='attention')
  thumbnail_path = os.path.join('thumbnails', name[:4], name[4:] + '.jpg')
  os.makedirs(os.path.dirname(thumbnail_path), exist_ok=True)
  thumb.write_to_file(thumbnail_path)

  # generate and save clip features
  features = extract_features(thumbnail_path)
  features_path = os.path.join('features', name[:4], name[4:] + '.bin')
  os.makedirs(os.path.dirname(features_path), exist_ok=True)
  with open(features_path, 'wb') as fp:
    vectorFormat = struct.Struct('f' * 512)
    fp.write(vectorFormat.pack(*features))

  # save to db
  cur.execute('INSERT INTO pictures VALUES (?, ?, ?, ?, ?, ?)', (path, date, image.width, image.height, thumbnail_path, features_path))
  con.commit()

 # process files recursively in directory
 def crawl(directory):
  for path in glob.iglob(directory + '/**/*.jpg', recursive=True):
    try:
      print(path)
      process(path)
    except Exception as e:
      print(e)

 # create an annoy index for quick vector search
 def create_index():
  index = AnnoyIndex(512, 'angular')
  for id, path in cur.execute('SELECT rowid, features FROM pictures'):
    vectorFormat = struct.Struct('f' * 512)
    with open(path, 'rb') as fp:
      features = vectorFormat.unpack(fp.read(2048))
    index.add_item(id, features)
  index.build(256)
  return index

 html = '''
 <!doctype html>
 <html lang="en-US">
 <head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width,initial-scale=1">
 </head>
 <body>
 %s
 </body>
 </html>'''

 # / redirects to /0
 @get('/')
 def callback():
  redirect('/0')

 # return 100 pictures with similar features to picture with given id
 @get('/<num:int>')
 def callback(num):
  result = '<form action="/search"><input type="text" name="q"><input type="submit"></form>'
  for i in index.get_nns_by_item(int(num), 100):
    path, thumbnail = cur.execute('SELECT path, thumbnail FROM pictures WHERE rowid = ?', (i,)).fetchone()
    result += '<a href="/%d"><img width="244" height="244" title="%s" src="%s"></a>' % (i, path, thumbnail)
  return html % result

 # convert text to vector with clip and return 100 similar pictures
 @get('/search')
 def callback():
  query = request.params.get('q')
  text = clip.tokenize([query]).cpu()
  text_onnx = text.detach().cpu().numpy().astype(np.int64)
  text_features = onnx_model.encode_text(text_onnx)[0]

  result = '<form action="/search"><input type="text" name="q" value="%s"><input type="submit"></form>' % query
  for i in index.get_nns_by_vector(text_features, 100):
    path, thumbnail = cur.execute('SELECT path, thumbnail FROM pictures WHERE rowid = ?', (i,)).fetchone()
    result += '<a href="/%d"><img width="244" height="244" title="%s" src="%s"></a>' % (i, path, thumbnail)
  return html % result

 # allow access to thumbnails
 @get("/thumbnails/<filepath:path>")
 def callback(filepath):
    return static_file(filepath, root="thumbnails")

 # generate thumbnails, features and index
 crawl(directory)
 index = create_index()

 # run server at http://localhost:8080/
 run(host='127.0.0.1', port=8080)
	# Searchable image gallery using CLIP vectors to represent picture content and text queries
	# * Install requirements:
	# pip install annoy pyvips git+https://github.com/Lednik7/CLIP-ONNX.git git+https://github.com/openai/CLIP.git bottle protobuf==3.20
	# * Download models:
	# wget https://clip-as-service.s3.us-east-2.amazonaws.com/models/onnx/ViT-B-32/{visual,textual}.onnx
	# * Run server and open http://127.0.0.1:8080:
	# python gallery.py /path/to/pictures/directory
	import os
	import sys
	import glob
	import struct
	import hashlib
	import sqlite3

	import clip
	import pyvips
	import numpy as np
	from PIL import Image
	from annoy import AnnoyIndex
	from clip_onnx import clip_onnx
	from bottle import run, get, static_file, redirect, request

	directory = sys.argv[1]

	os.makedirs('thumbnails', exist_ok=True)
	os.makedirs('features', exist_ok=True)

	# paths are stored in a simple sqlite db
	con = sqlite3.connect('db.sqlite')
	cur = con.cursor()
	cur.execute('CREATE TABLE IF NOT EXISTS pictures (path TEXT, date DATE, width INT, height INT, thumbnail TEXT, features TEXT)')
	con.commit()

	# load clip models
	_, preprocess = clip.load("ViT-B/32", device="cpu", jit=False)
	onnx_model = clip_onnx(None)
	onnx_model.load_onnx(visual_path="visual.onnx", textual_path="textual.onnx", logit_scale=100.0) # model.logit_scale.exp()
	onnx_model.start_sessions(providers=["CPUExecutionProvider"])

	# generate a vector of features with clip
	def extract_features(path):
	image_input = Image.open(path)
	image = preprocess(image_input).unsqueeze(0).cpu() # [1, 3, 224, 224]
	image_onnx = image.detach().cpu().numpy().astype(np.float32)
	image_features = onnx_model.encode_image(image_onnx)
	return image_features[0]

	# process a single picture
	def process(path):
	# pass already processed files
	if (1,) == cur.execute('SELECT 1 FROM pictures WHERE path = ?', (path,)).fetchone():
	return

	# compute unique identifier from path
	name = hashlib.sha1(path.encode('utf8')).hexdigest()

	# generate and save thumbnail
	image = pyvips.Image.new_from_file(path, access='sequential')
	image = image.autorot()
	date = image.get("exif-ifd0-DateTime")
	thumb = pyvips.Image.thumbnail_image(image, 224, crop='attention')
	thumbnail_path = os.path.join('thumbnails', name[:4], name[4:] + '.jpg')
	os.makedirs(os.path.dirname(thumbnail_path), exist_ok=True)
	thumb.write_to_file(thumbnail_path)

	# generate and save clip features
	features = extract_features(thumbnail_path)
	features_path = os.path.join('features', name[:4], name[4:] + '.bin')
	os.makedirs(os.path.dirname(features_path), exist_ok=True)
	with open(features_path, 'wb') as fp:
	vectorFormat = struct.Struct('f' * 512)
	fp.write(vectorFormat.pack(*features))

	# save to db
	cur.execute('INSERT INTO pictures VALUES (?, ?, ?, ?, ?, ?)', (path, date, image.width, image.height, thumbnail_path, features_path))
	con.commit()

	# process files recursively in directory
	def crawl(directory):
	for path in glob.iglob(directory + '/*/.jpg', recursive=True):
	try:
	print(path)
	process(path)
	except Exception as e:
	print(e)

	# create an annoy index for quick vector search
	def create_index():
	index = AnnoyIndex(512, 'angular')
	for id, path in cur.execute('SELECT rowid, features FROM pictures'):
	vectorFormat = struct.Struct('f' * 512)
	with open(path, 'rb') as fp:
	features = vectorFormat.unpack(fp.read(2048))
	index.add_item(id, features)
	index.build(256)
	return index

	html = '''
	<!doctype html>
	<html lang="en-US">
	<head>
	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width,initial-scale=1">
	</head>
	<body>
	%s
	</body>
	</html>'''

	# / redirects to /0
	@get('/')
	def callback():
	redirect('/0')

	# return 100 pictures with similar features to picture with given id
	@get('/<num:int>')
	def callback(num):
	result = '<form action="/search"><input type="text" name="q"><input type="submit"></form>'
	for i in index.get_nns_by_item(int(num), 100):
	path, thumbnail = cur.execute('SELECT path, thumbnail FROM pictures WHERE rowid = ?', (i,)).fetchone()
	result += '<a href="/%d"><img width="244" height="244" title="%s" src="%s"></a>' % (i, path, thumbnail)
	return html % result

	# convert text to vector with clip and return 100 similar pictures
	@get('/search')
	def callback():
	query = request.params.get('q')
	text = clip.tokenize([query]).cpu()
	text_onnx = text.detach().cpu().numpy().astype(np.int64)
	text_features = onnx_model.encode_text(text_onnx)[0]

	result = '<form action="/search"><input type="text" name="q" value="%s"><input type="submit"></form>' % query
	for i in index.get_nns_by_vector(text_features, 100):
	path, thumbnail = cur.execute('SELECT path, thumbnail FROM pictures WHERE rowid = ?', (i,)).fetchone()
	result += '<a href="/%d"><img width="244" height="244" title="%s" src="%s"></a>' % (i, path, thumbnail)
	return html % result

	# allow access to thumbnails
	@get("/thumbnails/<filepath:path>")
	def callback(filepath):
	return static_file(filepath, root="thumbnails")

	# generate thumbnails, features and index
	crawl(directory)
	index = create_index()

	# run server at http://localhost:8080/
	run(host='127.0.0.1', port=8080)