hjxdl 0.1.69__py3-none-any.whl → 0.1.71__py3-none-any.whl

hdl/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.69'
-__version_tuple__ = version_tuple = (0, 1, 69)
+__version__ = version = '0.1.71'
+__version_tuple__ = version_tuple = (0, 1, 71)

hdl/utils/llm/vis.py

@@ -11,6 +11,8 @@ import open_clip
 import natsort
 from redis.commands.search.field import VectorField
 from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from hdl.jupyfuncs.show.pbar import tqdm
+from redis.commands.search.query import Query
 
 from ..database_tools.connect import conn_redis
 
@@ -75,6 +77,7 @@ class ImgHandler:
         self.db_port = db_port
         self._db_conn = None
         self.num_vec_dim = num_vec_dim
+        self.pic_idx_name = "idx:pic_idx"
         if load_model:
             self.load_model()
 
@@ -142,10 +145,26 @@ class ImgHandler:
         Returns:
             torch.Tensor or numpy.ndarray: Image features extracted from the model.
         """
+
+        images_fixed = []
+        for img in images:
+            if isinstance(img, str):
+                if Path(Path(img).is_file()):
+                    images_fixed.append(Image.open(img))
+                elif img.startswith("data:image/jpeg;base64,"):
+                    images_fixed.append(imgbase64_to_pilimg(img))
+            elif isinstance(img, Image.Image):
+                images_fixed.append(img)
+            else:
+                raise TypeError(
+                    f"Not supported image type for {type(img)}"
+                )
+
+
         with torch.no_grad(), torch.amp.autocast("cuda"):
             imgs = torch.stack([
-                self.preprocess_val(Image.open(image)).to(self.device)
-                for image in images
+                self.preprocess_val(image).to(self.device)
+                for image in images_fixed
             ])
             img_features = self.model.encode_image(imgs, **kwargs)
             img_features /= img_features.norm(dim=-1, keepdim=True)
@@ -242,7 +261,11 @@ class ImgHandler:
             # > 0.9 很相似
             return sims
 
-    def vec_pics_todb(self, images):
+    def vec_pics_todb(
+        self,
+        images: list[str],
+        print_idx_info: bool = False
+    ):
         """Save image features to a Redis database.
 
         Args:
@@ -257,7 +280,7 @@ class ImgHandler:
         sorted_imgs = natsort.natsorted(images)
         img_feats = self.get_img_features(sorted_imgs, to_numpy=True)
         pipeline = self.db_conn.pipeline()
-        for img_file, emb in zip(sorted_imgs, img_feats):
+        for img_file, emb in tqdm(zip(sorted_imgs, img_feats)):
             # 初始化 Redis，先使用 img 文件名作为 Key 和 Value，后续再更新为图片特征向量
             # pipeline.json().set(img_file, "$", img_file)
             emb = emb.astype(np.float32).tolist()
@@ -281,31 +304,103 @@ class ImgHandler:
                 as_name="vector",  # 给这个字段定义一个别名，后续可以使用
             ),
         )
-        vector_idx_name = "idx:pic_idx"
+        # vector_idx_name = "idx:pic_idx"
         definition = IndexDefinition(
             prefix=["pic-"],
             index_type=IndexType.JSON
         )
         res = self.db_conn.ft(
-            vector_idx_name
+            self.pic_idx_name
         ).create_index(
             fields=schema,
             definition=definition
         )
         print("create_index:", res)
+        if print_idx_info:
+            print(self.pic_idx_info)
 
     @property
     def pic_idx_info(self):
         res = self.db_conn.ping()
         print("redis connected:", res)
-
-        vector_idx_name = "idx:pic_idx"
-
+        # vector_idx_name = "idx:pic_idx"
         # 从 Redis 数据库中读取索引状态
-        info = self.db_conn.ft(vector_idx_name).info()
+        info = self.db_conn.ft(self.pic_idx_name).info()
         # 获取索引状态中的 num_docs 和 hash_indexing_failures
         num_docs = info["num_docs"]
         indexing_failures = info["hash_indexing_failures"]
-        print(f"{num_docs} documents indexed with {indexing_failures} failures")
+        return (
+            f"{num_docs} documents indexed with {indexing_failures} failures"
+        )
+
+    def emb_search(
+        self,
+        emb_query,
+        num_max: int = 3,
+        extra_params: dict = None,
+    ):
+        """Search for similar embeddings in the database.
+
+        Args:
+            emb_query (str): The embedding query to search for.
+            num_max (int, optional): The maximum number of results to return. Defaults to 3.
+            extra_params (dict, optional): Extra parameters to include in the search query. Defaults to None.
+
+        Returns:
+            list: A list of tuples containing the document ID and JSON data for each result.
+        """
+        query = (
+            Query(
+                f"(*)=>[KNN {str(num_max)} @vector $query_vector AS vector_score]"
+            )
+            .sort_by("vector_score")
+            .return_fields("$")
+            .dialect(2)
+        )
+        if extra_params is None:
+            extra_params = {}
+        result_docs = (
+            self.db_conn.ft("idx:pic_idx")
+            .search(
+                query,
+                {
+                    "query_vector": emb_query
+                }
+                | extra_params,
+            )
+            .docs
+        )
+        results = [
+            (result_doc.id, json.loads(result_doc.json))
+            for result_doc in result_docs
+        ]
+        return results
+
+    def img_search(
+        self,
+        img,
+        num_max: int = 3,
+        extra_params: dict = None
+    ):
+        """Search for similar images in the database based on the input image.
+
+        Args:
+            img: Input image to search for similar images.
+            num_max: Maximum number of similar images to return (default is 3).
+            extra_params: Additional parameters to include in the search query (default is None).
+
+        Returns:
+            List of tuples containing the ID and JSON data of similar images found in the database.
+        """
+        emb_query = self.get_img_features(
+            [img], to_numpy=True
+        ).astype(np.float32)[0].tobytes()
+        results = self.emb_search(
+            emb_query=emb_query,
+            num_max=num_max,
+            extra_params=extra_params
+        )
+        return results
+
 
 

{hjxdl-0.1.69.dist-info → hjxdl-0.1.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.1.69
+Version: 0.1.71
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.1.69.dist-info → hjxdl-0.1.71.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=GffnD0jLJdhkd-vo989v40N90sQbofkayRBwxc6TVhQ,72
-hdl/_version.py,sha256=lVXaDORAcShTMS0Ix_T7zlcXAGst1lAMus_P8nq4GD4,413
+hdl/_version.py,sha256=Ov_2Aea_x3wREaTTkbdMMjDZcnW666PWCOUw0Yjj1mA,413
 hdl/args/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/args/loss_args.py,sha256=s7YzSdd7IjD24rZvvOrxLLFqMZQb9YylxKeyelSdrTk,70
 hdl/controllers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -131,12 +131,12 @@ hdl/utils/llm/chat.py,sha256=sk7Lw5Oa30k-l2fnJknkMmTc5zkBeEKsR981aeFhH5s,11907
 hdl/utils/llm/embs.py,sha256=Tf0FOYrOFZp7qQpEPiSCXzlgyHH0X9HVTUtsup74a9E,7174
 hdl/utils/llm/extract.py,sha256=2sK_WJzmYIc8iuWaM9DA6Nw3_6q1O4lJ5pKpcZo-bBA,6512
 hdl/utils/llm/llama_chat.py,sha256=watcHGOaz-bv3x-yDucYlGk5f8FiqfFhwWogrl334fk,4387
-hdl/utils/llm/vis.py,sha256=BygpjORSwbe_7SmThzt_OvzsrEym7FwCk6fubGyiDDc,10097
+hdl/utils/llm/vis.py,sha256=_xHEZ2_sRp7UiOWaI_6v96rBZwiQ7JkUF-DX5OHC7HE,13105
 hdl/utils/schedulers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/schedulers/norm_lr.py,sha256=bDwCmdEK-WkgxQMFBiMuchv8Mm7C0-GZJ6usm-PQk14,4461
 hdl/utils/weather/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/weather/weather.py,sha256=k11o6wM15kF8b9NMlEfrg68ak-SfSYLN3nOOflFUv-I,4381
-hjxdl-0.1.69.dist-info/METADATA,sha256=LZ8Y6-VsUCMYG1AQ5n-GBzff1lwfa-hOev1wT79OULc,903
-hjxdl-0.1.69.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-hjxdl-0.1.69.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.1.69.dist-info/RECORD,,
+hjxdl-0.1.71.dist-info/METADATA,sha256=jAEXEO3CVhotWgJjClsaKFlL59YBd2r6iotdUxo2o3g,903
+hjxdl-0.1.71.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+hjxdl-0.1.71.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.1.71.dist-info/RECORD,,