首页 > 数据库 > MySQL > TiDB Vector 太香啦:以图搜图初体验!
CREATE TABLE vector_table ( id int PRIMARY KEY, doc TEXT, embedding vector < float > (3) );
INSERT INTO vector_table VALUES (1, 'apple', '[1,1,1]'), (2, 'banana', '[1,1,2]'), (3, 'dog', '[2,2,2]');
SELECT *, vec_cosine_distance(embedding, '[1,1,3]') as distance FROM vector_table ORDER BY distance LIMIT 3;+-----------------------+-----------------------+---------------------+| id | doc | embedding | distance |+-----------------------+-----------------------+---------------------+| 2 | banana | [1,1,2] | 0.015268072165338209|| 3 | dog | [2,2,2] | 0.1296117202215108 || 1 | apple | [1,1,1] | 0.1296117202215108 |+---------+-------------+-----------------------+---------------------+
import pymysqldef GetConnection(): connection = pymysql.connect( host = "xxx.xxx.prod.aws.tidbcloud.com", port = 4000, user = "xxx.root", password = "xxx", database = "test", ssl_verify_cert = True, ssl_verify_identity = True, ssl_ca = "C:\\Users\\59131\\Downloads\\isrgrootx1.pem" ) return connection
from towhee import ops,pipe,AutoPipes,AutoConfig,DataCollectionimage_pipe = AutoPipes.pipeline('text_image_embedding')
@AutoConfig.registerclass TextImageEmbeddingConfig(BaseModel): model: Optional[str] = 'clip_vit_base_patch16' modality: Optional[str] = 'image' customize_embedding_op: Optional[Any] = None normalize_vec: Optional[bool] = True device: Optional[int] = -1
create table if not exists img_list ( id int PRIMARY KEY, path varchar(200) not null, embedding vector<float>(512));
def LoadImage(connection): cursor = connection.cursor() cursor.execute("create table if not exists img_list (id int PRIMARY KEY, path varchar(200) not null, embedding vector<float>(512));") img_dir='D:\\\\test\\\\' files = os.listdir(img_dir) for i in range(len(files)): path=os.path.join(img_dir, files[i]) embedding = image_pipe(path).get()[0] cursor.execute("INSERT INTO img_list VALUE ("+str(i)+",'"+path+"' , '"+np.array2string(embedding, separator=',')+"');") connection.commit()
def SearchInTiDB(connection,vector): cursor = connection.cursor() begin_time = datetime.datetime.now() cursor.execute("select id,path,vec_cosine_distance(embedding, '"+np.array2string(vector, separator=',')+"') as distance from img_list order by distance limit 3;") end_time=datetime.datetime.now() print("Search time:",(end_time-begin_time).total_seconds()) df =pd.DataFrame(cursor.fetchall()) return df[1]
def read_images(img_paths): imgs = [] op = ops.image_decode.cv2_rgb() for p in img_paths: imgs.append(op(p)) return imgs def ImageSearch(connection,path): emb = image_pipe(path).get()[0] res = SearchInTiDB(connection,emb) p = ( pipe.input('path','search_result') .map('path', 'img', ops.image_decode.cv2('rgb')) .map('search_result','prev',read_images) .output('img','prev') ) DataCollection(p(path,res)).show()
create table if not exists img_list_hnsw ( id int PRIMARY KEY, path varchar(200) not null, embedding vector<float>(512) COMMENT "hnsw(distance=cosine)");
E:\GitLocal\AITester>python tidb_vec.pySearch time: 0.320241+------------------------------------+------------------------------------------------------------------------------------------------------+| img | prev |+====================================+======================================================================================================+| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |+------------------------------------+------------------------------------------------------------------------------------------------------+E:\GitLocal\AITester>python tidb_vec.pySearch time: 0.239746+------------------------------------+------------------------------------------------------------------------------------------------------+| img | prev |+====================================+======================================================================================================+| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |+------------------------------------+------------------------------------------------------------------------------------------------------+
def TextSearch(connection,text): text_conf = AutoConfig.load_config('text_image_embedding') text_conf.modality = 'text' text_pipe = AutoPipes.pipeline('text_image_embedding', text_conf) embedding = text_pipe(text).get()[0] res=SearchInTiDB(connection,embedding) p = ( pipe.input('text','search_result') .map('search_result','prev',read_images) .output('text','prev') ) DataCollection(p(text,res)).show()
来自 “ ITPUB博客 ” ,链接:https://blog.itpub.net/69994146/viewspace-3014307/ ,如需转载,请注明出处,否则将追究法律责任。