hjxdl 0.1.61__py3-none-any.whl → 0.1.63__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.61'
-__version_tuple__ = version_tuple = (0, 1, 61)
+__version__ = version = '0.1.63'
+__version_tuple__ = version_tuple = (0, 1, 63)

hdl/utils/llm/vis.py

@@ -66,6 +66,11 @@ class ImgHandler:
 
     @property
     def redis_conn(self):
+        """Establishes a connection to Redis server if not already connected.
+
+            Returns:
+                Redis connection object: A connection to the Redis server.
+        """
         if self._redis_conn is None:
             self._redis_conn = conn_redis(self.redis_host, self.redis_port)
         return self._redis_conn
@@ -76,14 +81,25 @@ class ImgHandler:
         to_numpy = False,
         **kwargs
     ):
-        imgs = torch.stack([
-            self.preprocess_val(Image.open(image)).to(self.device)
-            for image in images
-        ])
-        img_features = self.model.encode_image(imgs, **kwargs)
-        img_features /= img_features.norm(dim=-1, keepdim=True)
-        if to_numpy:
-            img_features = img_features.cpu().numpy()
+        """Get image features using a pretrained model.
+
+        Args:
+            images (list): List of image paths.
+            to_numpy (bool, optional): Whether to convert the image features to numpy array. Defaults to False.
+            **kwargs: Additional keyword arguments to pass to the model.
+
+        Returns:
+            torch.Tensor or numpy.ndarray: Image features extracted from the model.
+        """
+        with torch.no_grad(), torch.amp.autocast("cuda"):
+            imgs = torch.stack([
+                self.preprocess_val(Image.open(image)).to(self.device)
+                for image in images
+            ])
+            img_features = self.model.encode_image(imgs, **kwargs)
+            img_features /= img_features.norm(dim=-1, keepdim=True)
+            if to_numpy:
+                img_features = img_features.cpu().numpy()
         return img_features
 
     def get_text_features(
@@ -92,14 +108,28 @@ class ImgHandler:
         to_numpy = False,
         **kwargs
     ):
-        txts = self.tokenizer(
-            texts,
-            context_length=self.model.context_length
-        ).to(self.device)
-        txt_features = self.model.encode_text(txts, **kwargs)
-        txt_features /= txt_features.norm(dim=-1, keepdim=True)
-        if to_numpy:
-            txt_features = txt_features.cpu().numpy()
+        """Get text features from the input texts.
+
+        Args:
+            texts (list): List of input texts to extract features from.
+            to_numpy (bool, optional): Whether to convert the output features to a numpy array. Defaults to False.
+            **kwargs: Additional keyword arguments to pass to the model for encoding text.
+
+        Returns:
+            torch.Tensor or numpy.ndarray: Extracted text features.
+
+        Example:
+            get_text_features(["text1", "text2"], to_numpy=True)
+        """
+        with torch.no_grad(), torch.amp.autocast("cuda"):
+            txts = self.tokenizer(
+                texts,
+                context_length=self.model.context_length
+            ).to(self.device)
+            txt_features = self.model.encode_text(txts, **kwargs)
+            txt_features /= txt_features.norm(dim=-1, keepdim=True)
+            if to_numpy:
+                txt_features = txt_features.cpu().numpy()
         return txt_features
 
 
@@ -111,13 +141,54 @@ class ImgHandler:
         to_numpy: bool = False,
         **kwargs
     ):
-        image_features = self.get_img_features(images, **kwargs)
-        text_features = self.get_text_features(texts, **kwargs)
-        text_probs = (100.0 * image_features @ text_features.T)
-        if probs:
-            text_probs = text_probs.softmax(dim=-1)
-        if to_numpy:
-            text_probs = text_probs.cpu().numpy()
+        """Get the probabilities of text-image associations.
+
+        Args:
+            texts (list): List of text inputs.
+            images (list): List of image inputs.
+            probs (bool, optional): Whether to return probabilities. Defaults to False.
+            to_numpy (bool, optional): Whether to convert the output to a numpy array. Defaults to False.
+            **kwargs: Additional keyword arguments for feature extraction.
+
+        Returns:
+            torch.Tensor or numpy.ndarray: Text-image association probabilities.
+        """
+        with torch.no_grad(), torch.amp.autocast("cuda"):
+            image_features = self.get_img_features(images, **kwargs)
+            text_features = self.get_text_features(texts, **kwargs)
+            text_probs = (100.0 * image_features @ text_features.T)
+            # >3 有关联
+            if probs:
+                text_probs = text_probs.softmax(dim=-1)
+            if to_numpy:
+                text_probs = text_probs.cpu().numpy()
         return text_probs
 
+    def get_pics_sims(
+        self,
+        images1,
+        images2,
+        to_numpy: bool = False,
+        **kwargs
+    ):
+        """Calculate similarity scores between two sets of images.
+
+            Args:
+                images1: First set of images.
+                images2: Second set of images.
+                to_numpy: Whether to convert the similarity scores to a numpy array (default is False).
+                **kwargs: Additional keyword arguments to pass to the get_img_features method.
+
+            Returns:
+                torch.Tensor or numpy.ndarray: Similarity scores between the two sets of images.
+        """
+        with torch.no_grad(), torch.amp.autocast("cuda"):
+            img1_feats = self.get_img_features(images1, **kwargs)
+            img2_feats = self.get_img_features(images2, **kwargs)
+            sims = img1_feats @ img2_feats.T
+            if to_numpy:
+                sims = sims.cpu().numpy()
+            # > 0.9 很相似
+            return sims
+
 

{hjxdl-0.1.61.dist-info → hjxdl-0.1.63.dist-info}/METADATA

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

{hjxdl-0.1.61.dist-info → hjxdl-0.1.63.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=GffnD0jLJdhkd-vo989v40N90sQbofkayRBwxc6TVhQ,72
-hdl/_version.py,sha256=iuq68v789XjbJHhRQlyaD2e2yTTa7FKtdYjoKlGAJmY,413
+hdl/_version.py,sha256=mjLtl9xPbny0ZTEF6CvbbjoK1aXIg4Io5ZPI-Nfqo_k,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=zVgy3MepIEKk9oeXuADjloBaWjlshjZB6g2uHbtcPEg,3471
+hdl/utils/llm/vis.py,sha256=6CpV9lXvHw5TRSsu2XKLDNYJIP3QsKYEDVY7pKx9r1A,6439
 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.61.dist-info/METADATA,sha256=Ud4NYOjhC49-QV63fDwSy0AltmVxPtFz_HcEr5qi9Bs,880
-hjxdl-0.1.61.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-hjxdl-0.1.61.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.1.61.dist-info/RECORD,,
+hjxdl-0.1.63.dist-info/METADATA,sha256=LzK9HNry28LVlCSbab50Gp1hoIw-vQFDuCzwCO1Zyjk,880
+hjxdl-0.1.63.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+hjxdl-0.1.63.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.1.63.dist-info/RECORD,,