ultralytics 8.2.83__py3-none-any.whl → 8.2.84__py3-none-any.whl

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.2.83"
+__version__ = "8.2.84"
 
 import os
 

ultralytics/engine/model.py

@@ -128,8 +128,10 @@ class Model(nn.Module):
         if self.is_hub_model(model):
             # Fetch model from HUB
             checks.check_requirements("hub-sdk>=0.0.8")
-            self.session = HUBTrainingSession.create_session(model)
-            model = self.session.model_file
+            session = HUBTrainingSession.create_session(model)
+            model = session.model_file
+            if session.train_args:  # training sent from HUB
+                self.session = session
 
         # Check if Triton Server model
         elif self.is_triton_model(model):

ultralytics/hub/session.py

@@ -50,6 +50,7 @@ class HUBTrainingSession:
         self.model = None
         self.model_url = None
         self.model_file = None
+        self.train_args = None
 
         # Parse input
         api_key, model_id, self.filename = self._parse_identifier(identifier)

ultralytics/models/sam/model.py

@@ -106,7 +106,7 @@ class SAM(Model):
             ...     print(f"Detected {len(r.masks)} masks")
         """
         overrides = dict(conf=0.25, task="segment", mode="predict", imgsz=1024)
-        kwargs.update(overrides)
+        kwargs = {**overrides, **kwargs}
         prompts = dict(bboxes=bboxes, points=points, labels=labels)
         return super().predict(source, stream, prompts=prompts, **kwargs)
 

ultralytics/models/sam/modules/encoders.py

@@ -151,7 +151,12 @@ class ImageEncoderViT(nn.Module):
         """Processes input through patch embedding, positional embedding, transformer blocks, and neck module."""
         x = self.patch_embed(x)
         if self.pos_embed is not None:
-            x = x + self.pos_embed
+            pos_embed = (
+                F.interpolate(self.pos_embed.permute(0, 3, 1, 2), scale_factor=self.img_size / 1024).permute(0, 2, 3, 1)
+                if self.img_size != 1024
+                else self.pos_embed
+            )
+            x = x + pos_embed
         for blk in self.blocks:
             x = blk(x)
         return self.neck(x.permute(0, 3, 1, 2))

ultralytics/models/sam/modules/sam.py

@@ -90,6 +90,19 @@ class SAMModel(nn.Module):
         self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
         self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
 
+    def set_imgsz(self, imgsz):
+        """
+        Set image size to make model compatible with different image sizes.
+
+        Args:
+            imgsz (Tuple[int, int]): The size of the input image.
+        """
+        if hasattr(self.image_encoder, "set_imgsz"):
+            self.image_encoder.set_imgsz(imgsz)
+        self.prompt_encoder.input_image_size = imgsz
+        self.prompt_encoder.image_embedding_size = [x // 16 for x in imgsz]  # 16 is fixed as patch size of ViT model
+        self.image_encoder.img_size = imgsz[0]
+
 
 class SAM2Model(torch.nn.Module):
     """
@@ -940,3 +953,14 @@ class SAM2Model(torch.nn.Module):
         # don't overlap (here sigmoid(-10.0)=4.5398e-05)
         pred_masks = torch.where(keep, pred_masks, torch.clamp(pred_masks, max=-10.0))
         return pred_masks
+
+    def set_imgsz(self, imgsz):
+        """
+        Set image size to make model compatible with different image sizes.
+
+        Args:
+            imgsz (Tuple[int, int]): The size of the input image.
+        """
+        self.image_size = imgsz[0]
+        self.sam_prompt_encoder.input_image_size = imgsz
+        self.sam_prompt_encoder.image_embedding_size = [x // 16 for x in imgsz]  # fixed ViT patch size of 16

ultralytics/models/sam/modules/tiny_encoder.py

@@ -982,10 +982,31 @@ class TinyViT(nn.Module):
             layer = self.layers[i]
             x = layer(x)
         batch, _, channel = x.shape
-        x = x.view(batch, 64, 64, channel)
+        x = x.view(batch, self.patches_resolution[0] // 4, self.patches_resolution[1] // 4, channel)
         x = x.permute(0, 3, 1, 2)
         return self.neck(x)
 
     def forward(self, x):
         """Performs the forward pass through the TinyViT model, extracting features from the input image."""
         return self.forward_features(x)
+
+    def set_imgsz(self, imgsz=[1024, 1024]):
+        """
+        Set image size to make model compatible with different image sizes.
+
+        Args:
+            imgsz (Tuple[int, int]): The size of the input image.
+        """
+        imgsz = [s // 4 for s in imgsz]
+        self.patches_resolution = imgsz
+        for i, layer in enumerate(self.layers):
+            input_resolution = (
+                imgsz[0] // (2 ** (i - 1 if i == 3 else i)),
+                imgsz[1] // (2 ** (i - 1 if i == 3 else i)),
+            )
+            layer.input_resolution = input_resolution
+            if layer.downsample is not None:
+                layer.downsample.input_resolution = input_resolution
+            if isinstance(layer, BasicLayer):
+                for b in layer.blocks:
+                    b.input_resolution = input_resolution

ultralytics/models/sam/predict.py

@@ -95,7 +95,7 @@ class Predictor(BasePredictor):
         """
         if overrides is None:
             overrides = {}
-        overrides.update(dict(task="segment", mode="predict", imgsz=1024))
+        overrides.update(dict(task="segment", mode="predict"))
         super().__init__(cfg, overrides, _callbacks)
         self.args.retina_masks = True
         self.im = None
@@ -455,8 +455,11 @@ class Predictor(BasePredictor):
                 cls = torch.arange(len(pred_masks), dtype=torch.int32, device=pred_masks.device)
                 pred_bboxes = torch.cat([pred_bboxes, pred_scores[:, None], cls[:, None]], dim=-1)
 
-            masks = ops.scale_masks(masks[None].float(), orig_img.shape[:2], padding=False)[0]
-            masks = masks > self.model.mask_threshold  # to bool
+            if len(masks) == 0:
+                masks = None
+            else:
+                masks = ops.scale_masks(masks[None].float(), orig_img.shape[:2], padding=False)[0]
+                masks = masks > self.model.mask_threshold  # to bool
             results.append(Results(orig_img, path=img_path, names=names, masks=masks, boxes=pred_bboxes))
         # Reset segment-all mode.
         self.segment_all = False
@@ -522,6 +525,10 @@ class Predictor(BasePredictor):
 
     def get_im_features(self, im):
         """Extracts image features using the SAM model's image encoder for subsequent mask prediction."""
+        assert (
+            isinstance(self.imgsz, (tuple, list)) and self.imgsz[0] == self.imgsz[1]
+        ), f"SAM models only support square image size, but got {self.imgsz}."
+        self.model.set_imgsz(self.imgsz)
         return self.model.image_encoder(im)
 
     def set_prompts(self, prompts):
@@ -761,6 +768,12 @@ class SAM2Predictor(Predictor):
 
     def get_im_features(self, im):
         """Extracts image features from the SAM image encoder for subsequent processing."""
+        assert (
+            isinstance(self.imgsz, (tuple, list)) and self.imgsz[0] == self.imgsz[1]
+        ), f"SAM 2 models only support square image size, but got {self.imgsz}."
+        self.model.set_imgsz(self.imgsz)
+        self._bb_feat_sizes = [[x // (4 * i) for x in self.imgsz] for i in [1, 2, 4]]
+
         backbone_out = self.model.forward_image(im)
         _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)
         if self.model.directly_add_no_mem_embed:

{ultralytics-8.2.83.dist-info → ultralytics-8.2.84.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.2.83
+Version: 8.2.84
 Summary: Ultralytics YOLOv8 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author: Glenn Jocher, Ayush Chaurasia, Jing Qiu
 Maintainer: Glenn Jocher, Ayush Chaurasia, Jing Qiu

{ultralytics-8.2.83.dist-info → ultralytics-8.2.84.dist-info}/RECORD

@@ -8,7 +8,7 @@ tests/test_exports.py,sha256=Uezf3OatpPHlo5qoPw-2kqkZxuMCF9L4XF2riD4vmII,8225
 tests/test_integrations.py,sha256=xglcfMPjfVh346PV8WTpk6tBxraCXEFJEQyyJMr5tyU,6064
 tests/test_python.py,sha256=SxBf5GNu7vXQP8QxTlSOzCzcQNN0PLA6EX8M33VDHsU,21927
 tests/test_solutions.py,sha256=p_2edhl96Ty3jwzSf02Q2m2mTu9skc0Z-eMcUuuXfLg,3300
-ultralytics/__init__.py,sha256=GCuzRCtg8kNcfqSooPJeB90ILBmOEWf_LhJirWsy0Z4,694
+ultralytics/__init__.py,sha256=YbOKlDA48wG1Ggkk8XH3ZxSQHQ2ys3jxPb9nb-H0lhk,694
 ultralytics/assets/bus.jpg,sha256=wCAZxJecGR63Od3ZRERe9Aja1Weayrb9Ug751DS_vGM,137419
 ultralytics/assets/zidane.jpg,sha256=Ftc4aeMmen1O0A3o6GCDO9FlfBslLpTAw0gnetx7bts,50427
 ultralytics/cfg/__init__.py,sha256=pkB7wk0pHOA3xzKzMbS-hA0iJoPOWVNnwZJh0LuWh-w,33089
@@ -99,7 +99,7 @@ ultralytics/data/explorer/gui/__init__.py,sha256=mHtJuK4hwF8cuV-VHDc7tp6u6D1gHz2
 ultralytics/data/explorer/gui/dash.py,sha256=vZ476NaUH4FKU08rAJ1K9WNyKtg0soMyJJxqg176yWc,10498
 ultralytics/engine/__init__.py,sha256=mHtJuK4hwF8cuV-VHDc7tp6u6D1gHz2Z7JI8grmQDTs,42
 ultralytics/engine/exporter.py,sha256=R46KseDRaSoPo8NTZX2yybxjJoLP8KCXIgMHNagE314,58888
-ultralytics/engine/model.py,sha256=2XsZYBZF3O4LEqRyffTudP31x7KmyV7jXzoIEMu8Mp8,52191
+ultralytics/engine/model.py,sha256=AB9tu7kJW-QiTAp0F_J8KQJ4FijsHXcYBTaVHb7aMrg,52281
 ultralytics/engine/predictor.py,sha256=TtCOhjCOAz4iUXog8V2h_3VhsKPJM8HmP3i6W_qqdhk,17460
 ultralytics/engine/results.py,sha256=PgRcz90S7eMwlogqEvax8O1sU3CPA2tEmrAL5kSr6II,70537
 ultralytics/engine/trainer.py,sha256=ebFsES6KfVlVoCx9xeEpGDtVDumEndTHqojbcs9BzHg,35940
@@ -107,7 +107,7 @@ ultralytics/engine/tuner.py,sha256=gPqDTHH7vRB2O3YyH26m1BjVKbXxuA2XAlPRzTKFZsc,1
 ultralytics/engine/validator.py,sha256=u7qh9AiHasfhIqwojfjvAY8B2IIhd2MIHRwTxsTP6RY,14586
 ultralytics/hub/__init__.py,sha256=AM_twjV9ouUmyxh3opoPgTqDpMOd8xIOHsAKdWS2L18,5663
 ultralytics/hub/auth.py,sha256=kDLakGa2NbzvMAeXc2UdzZ65r0AH-XeM_JfsDY97WGk,5545
-ultralytics/hub/session.py,sha256=JgiCJ7lre6wW6B5Xylo_-f7x0_QkAAF-xCbTVl1eank,16857
+ultralytics/hub/session.py,sha256=_5yQNKkeaOnxwBeL85ueCgR-IYnDQ89WuzFNjTNPflU,16888
 ultralytics/hub/utils.py,sha256=tXfM3QbXBcf4Y6StgHI1pktT4OM7Ic9eF3xiBFHGlhY,9721
 ultralytics/hub/google/__init__.py,sha256=qyvvpGP-4NAtrn7GLqfqxP_aWuRP1T0OvJYafWKvL2Q,7512
 ultralytics/models/__init__.py,sha256=TT9iLCL_n9Y80dcUq0Fo-p-GRZCSU2vrWXM3CoMwqqE,265
@@ -128,15 +128,15 @@ ultralytics/models/rtdetr/val.py,sha256=xVjZShZ1AvES97wVekl2q_1g20Pq-IIHhkJdWtxM
 ultralytics/models/sam/__init__.py,sha256=o4_D6y8YJlOXIK7Lwo9RHnIJJ9xoFNi4zK99QSc1kdM,176
 ultralytics/models/sam/amg.py,sha256=GrmO_8YfIDt_QkPEMF_WFjPZkhwhf7iwx7ig8JgOUnE,8709
 ultralytics/models/sam/build.py,sha256=zNQbrgSHUgz1gyXQwLKGTpa6CSEjeaevcP3w1Z1l3mo,12233
-ultralytics/models/sam/model.py,sha256=_EAgXA7nLeG_-wnvgG3tM_V4oDKlsHGTWX8El4xomo4,7374
-ultralytics/models/sam/predict.py,sha256=CFZriLbrH_ZoxP7DyYESuqw0rIrehu8Hctzr-IKucb0,37744
+ultralytics/models/sam/model.py,sha256=2KFUp8SHiqOgwUjkdqdau0oduJwKQxm4N9GHWjdhUFo,7382
+ultralytics/models/sam/predict.py,sha256=4HOvBp27MvO8ef3gD64wVooNT1P5eMy3Bk8W7ysU57o,38352
 ultralytics/models/sam/modules/__init__.py,sha256=mHtJuK4hwF8cuV-VHDc7tp6u6D1gHz2Z7JI8grmQDTs,42
 ultralytics/models/sam/modules/blocks.py,sha256=qXCXMqkQG0fpAvCkA9TrtimfOLDtyJfCx3bDfh3bJUs,45974
 ultralytics/models/sam/modules/decoders.py,sha256=qDr12mDvDA-VIMI7Q9oIYBG9DQcvxDFpPzyAjyqrcbg,25896
-ultralytics/models/sam/modules/encoders.py,sha256=KvQFAtqfGvCAr4kcMXxnJvjwIhaQ0a3Wwp0KhSSG_oA,34615
+ultralytics/models/sam/modules/encoders.py,sha256=vDOv8sdbcWc31aVn7hg-JyLP6CRziPep5FPDG2wxwns,34848
 ultralytics/models/sam/modules/memory_attention.py,sha256=XilWBnRfH8wZxIoL2-yEk-dRypCsS0Jf_9t8WJxXKg0,9722
-ultralytics/models/sam/modules/sam.py,sha256=6GE0E4m1J91QgFeUo0vw-Cz7gxebD1VfIGTNKm_meUI,49558
-ultralytics/models/sam/modules/tiny_encoder.py,sha256=04btkm8KfLZBP-nPihFmpO-mNrD2euVFSZtuzc21IZk,40439
+ultralytics/models/sam/modules/sam.py,sha256=0Df9psft2-uShp-WTP1oZT6x5QSE9S0i7XKBdZ4tpfE,50507
+ultralytics/models/sam/modules/tiny_encoder.py,sha256=NyzeFMLnmqwcFQFs-JBM9PCWSsYoYZ_6h59Un1DeDV0,41332
 ultralytics/models/sam/modules/transformer.py,sha256=oMlns0i_bcEqdcdnDJzeM7er2_yqqdYk4hZd3QbEGWQ,16154
 ultralytics/models/sam/modules/utils.py,sha256=Y36V6BVy6GeaAvKE8gHmoDIa-f5LjJpmSVwywNkv2yk,12315
 ultralytics/models/utils/__init__.py,sha256=mHtJuK4hwF8cuV-VHDc7tp6u6D1gHz2Z7JI8grmQDTs,42
@@ -225,9 +225,9 @@ ultralytics/utils/callbacks/neptune.py,sha256=5Z3ua5YBTUS56FH8VQKQG1aaIo9fH8GEyz
 ultralytics/utils/callbacks/raytune.py,sha256=ODVYzy-CoM4Uge0zjkh3Hnh9nF2M0vhDrSenXnvcizw,705
 ultralytics/utils/callbacks/tensorboard.py,sha256=0kn4IR10no99UCIheojWRujgybmUHSx5fPI6Vsq6l_g,4135
 ultralytics/utils/callbacks/wb.py,sha256=9-fjQIdLjr3b73DTE3rHO171KvbH1VweJ-bmbv-rqTw,6747
-ultralytics-8.2.83.dist-info/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
-ultralytics-8.2.83.dist-info/METADATA,sha256=bcRaFabp1etZymVCQyCNkWtJeVceTpSGXSOmywZaBiA,41778
-ultralytics-8.2.83.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
-ultralytics-8.2.83.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
-ultralytics-8.2.83.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
-ultralytics-8.2.83.dist-info/RECORD,,
+ultralytics-8.2.84.dist-info/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
+ultralytics-8.2.84.dist-info/METADATA,sha256=NPkRQek_u6FUwJA7J-GnbT4OSmqxRJYZTR6CIWVTE24,41778
+ultralytics-8.2.84.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+ultralytics-8.2.84.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
+ultralytics-8.2.84.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
+ultralytics-8.2.84.dist-info/RECORD,,