♻️ [Refactor] move to_tensor to transform

yolo/__init__.py

@@ -0,0 +1,23 @@
+from yolo.config.config import Config
+from yolo.model.yolo import create_model
+from yolo.tools.data_loader import AugmentationComposer, create_dataloader
+from yolo.tools.drawer import draw_bboxes
+from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.utils.bounding_box_utils import bbox_nms
+from yolo.utils.deploy_utils import FastModelLoader
+from yolo.utils.logging_utils import custom_logger
+
+all = [
+    "create_model",
+    "Config",
+    "custom_logger",
+    "validate_log_directory",
+    "draw_bboxes",
+    "bbox_nms",
+    "AugmentationComposer",
+    "create_dataloader",
+    "FastModelLoader",
+    "ModelTester",
+    "ModelTrainer",
+    "ModelValidator",
+]

yolo/tools/data_augmentation.py

@@ -21,6 +21,7 @@ class AugmentationComposer:
         for transform in self.transforms:
             image, boxes = transform(image, boxes)
         image, boxes = self.pad_resize(image, boxes)
+        image = TF.to_tensor(image)
         return image, boxes
 
 

yolo/tools/data_loader.py

@@ -151,9 +151,7 @@ class YoloDataset(Dataset):
 
     def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]:
         img, bboxes = self.get_data(idx)
-        if self.transform:
-            img, bboxes = self.transform(img, bboxes)
-        img = TF.to_tensor(img)
+        img, bboxes = self.transform(img, bboxes)
         return img, bboxes
 
     def __len__(self) -> int:
@@ -269,7 +267,7 @@ class StreamDataLoader:
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             frame = Image.fromarray(frame)
         frame, _ = self.transform(frame, torch.zeros(0, 5))
-        frame = TF.to_tensor(frame)[None]
+        frame = frame[None]
         if not self.is_stream:
             self.queue.put(frame)
         else: