henry000 commited on
Commit
5e0d5e2
·
1 Parent(s): 2522f72

🔨 [Add] image ratio min-max constrain

Browse files
Files changed (1) hide show
  1. yolo/tools/data_loader.py +2 -3
yolo/tools/data_loader.py CHANGED
@@ -141,7 +141,7 @@ class YoloDataset(Dataset):
141
  if bboxes:
142
  return torch.stack(bboxes)
143
  else:
144
- logger.warning("No valid BBox in {}", label_path)
145
  return torch.zeros((0, 5))
146
 
147
  def get_data(self, idx):
@@ -158,8 +158,7 @@ class YoloDataset(Dataset):
158
  def _update_image_size(self, idx: int) -> None:
159
  """Update image size based on dynamic shape and batch settings."""
160
  batch_start_idx = (idx // self.batch_size) * self.batch_size
161
- image_ratio = self.ratios[batch_start_idx]
162
-
163
  shift = ((self.base_size / 32 * (image_ratio - 1)) // (image_ratio + 1)) * 32
164
 
165
  self.image_size = [int(self.base_size + shift), int(self.base_size - shift)]
 
141
  if bboxes:
142
  return torch.stack(bboxes)
143
  else:
144
+ logger.warning(f"No valid BBox in {label_path}")
145
  return torch.zeros((0, 5))
146
 
147
  def get_data(self, idx):
 
158
  def _update_image_size(self, idx: int) -> None:
159
  """Update image size based on dynamic shape and batch settings."""
160
  batch_start_idx = (idx // self.batch_size) * self.batch_size
161
+ image_ratio = self.ratios[batch_start_idx].clip(1 / 3, 3)
 
162
  shift = ((self.base_size / 32 * (image_ratio - 1)) // (image_ratio + 1)) * 32
163
 
164
  self.image_size = [int(self.base_size + shift), int(self.base_size - shift)]