feat: enhance image size handling across the codebase

README.md

@@ -69,8 +69,8 @@ layout = DatasetLayout.ROBOFLOW_COCO
 auto_dataset = AutoDataset(dataset_name=ds_name, task=task, layout=layout)
 
 train_augs, val_augs = get_default_by_task(task, 640, advanced=False)
-train_dataset = auto_dataset.get_split(augs=train_augs.get_augmentations(), split=DatasetSplitType.TRAIN)
-valid_dataset = auto_dataset.get_split(augs=val_augs.get_augmentations(), split=DatasetSplitType.VAL)
+train_dataset = auto_dataset.get_split(augs=train_augs, split=DatasetSplitType.TRAIN)
+valid_dataset = auto_dataset.get_split(augs=val_augs, split=DatasetSplitType.VAL)
 
 
 model = ModelManager.get("fai-detr-l-obj365")

docs/cli.md

@@ -85,9 +85,12 @@ Where:
 ### 🏋️ Training
 
 ```bash
-# Basic training
+# Basic training (square image)
 focoos train --model fai-detr-m-coco --dataset mydataset.zip --im-size 640
 
+# Training with non-square resolution
+focoos train --model fai-detr-m-coco --dataset mydataset.zip --im-size 640,480
+
 # Advanced training with custom hyperparameters
 focoos train \
   --model fai-detr-m-coco \
@@ -157,9 +160,12 @@ focoos predict \
 ### 📤 Model Export
 
 ```bash
-# Export to ONNX (default)
+# Export to ONNX (default, square image)
 focoos export --model fai-detr-m-coco --im-size 640
 
+# Export to ONNX with non-square resolution
+focoos export --model fai-detr-m-coco --im-size 640,480
+
 # Export to TorchScript
 focoos export \
   --model fai-detr-m-coco \
@@ -249,7 +255,7 @@ The interface will automatically open in your default web browser, typically at
 | `--model`      | Model name or path          | **Required**   | `fai-detr-m-coco`, `path/to/model` |
 | `--dataset`    | Dataset name or path        | **Required**   | `mydataset.zip`, `path/to/data/`   |
 | `--source`     | Input source (predict only) | **Required**   | `image.jpg`                        |
-| `--im-size`    | Input image size            | 640            | Any positive integer               |
+| `--im-size`    | Input image size            | 640            | Integer (square) or "height,width" or "heightxwidth" (non-square) |
 | `--batch-size` | Batch size                  | 16             | Powers of 2 recommended            |
 | `--device`     | Compute device              | `cuda`         | `cuda`, `cpu`               |
 | `--workers`    | Data loading workers        | 4              | 0-16 recommended                   |
@@ -319,7 +325,7 @@ Use CLI commands programmatically in Python:
 ```python
 from focoos.cli.commands import train_command, predict_command, export_command
 
-# Train a model
+# Train a model (square image)
 train_command(
     model_name="fai-detr-m-coco",
     dataset_name="mydataset.zip",
@@ -330,6 +336,17 @@ train_command(
     batch_size=16
 )
 
+# Train a model with non-square resolution
+train_command(
+    model_name="fai-detr-m-coco",
+    dataset_name="mydataset.zip",
+    dataset_layout="roboflow_coco",
+    im_size=(640, 480),  # (height, width)
+    run_name="my_training",
+    max_iters=5000,
+    batch_size=16
+)
+
 # Run inference
 results = predict_command(
     model_name="fai-detr-m-coco",
@@ -364,8 +381,11 @@ focoos version
 # Reduce batch size
 focoos train --model fai-detr-m-coco --dataset data.zip --batch-size 8
 
-# Use smaller image size
+# Use smaller image size (square)
 focoos train --model fai-detr-m-coco --dataset data.zip --im-size 480
+
+# Use non-square image size for memory efficiency
+focoos train --model fai-detr-m-coco --dataset data.zip --im-size 480,360
 ```
 
 **Dataset not found:**

docs/concepts.md

@@ -120,7 +120,7 @@ model_info = ModelInfo(
     name="custom_detector",
     model_family=ModelFamily.DETR,
     classes=["person", "car", "bicycle"],
-    im_size=640,
+    im_size=640,  # Square image (640x640), or use (640, 480) for non-square (height, width)
     task=Task.DETECTION,
     config={
         "num_classes": 3,
@@ -220,7 +220,7 @@ model.train(train_args, train_dataset, val_dataset, hub=hub)
 Exports the model to different runtime formats for optimized inference. The main function arguments are:
  - `runtime_type`: specify the target runtime and must be one of the supported (see [RuntimeType](/focoos/api/ports/#focoos.ports.RuntimeType))
  - `out_dir`: the destination folder for the exported model
- - `image_size`: the target image size, as an optional integer
+ - `image_size`: the target image size, as an optional integer (square) or tuple (height, width) for non-square images
 
 The function returns an [`InferModel`](#infer-model) instance for the exported model.
 

docs/inference.md

@@ -151,7 +151,11 @@ runtime = RuntimeType.TORCHSCRIPT_32
 It's time to export the model. We can use the export method of the models.
 
 ```python
+# Export with square resolution (512x512)
 optimized_model = model.export(runtime_type=runtime, image_size=512)
+
+# Export with non-square resolution (640x480, height x width)
+optimized_model = model.export(runtime_type=runtime, image_size=(640, 480))
 ```
 
 Let's visualize the output. As you will see, there are not differences from the model in pure torch.
@@ -183,5 +187,9 @@ optimized_model = model.export(runtime_type=runtime)
 detections = optimized_model(image)
 display(annotate_image(image, detections, task=model.model_info.task, classes=model.model_info.classes))
 
+# Benchmark with square resolution
 optimized_model.benchmark(iterations=10, size=640)
+
+# Benchmark with non-square resolution
+optimized_model.benchmark(iterations=10, size=(640, 480))
 ```

docs/training.md

@@ -65,7 +65,11 @@ task = dataset.task  # see ports.Task for more information
 layout = dataset.layout  # see ports.DatasetLayout for more information
 auto_dataset = AutoDataset(dataset_name=dataset_path, task=task, layout=layout)
 
-augs = DatasetAugmentations(resolution=512).get_augmentations()
+# Square resolution (512x512)
+augs = DatasetAugmentations(resolution=512)
+
+# Non-square resolution (640x480, height x width)
+augs = DatasetAugmentations(resolution=(640, 480))
 
 train_dataset = auto_dataset.get_split(augs=augs, split=DatasetSplitType.TRAIN)
 valid_dataset = auto_dataset.get_split(augs=augs, split=DatasetSplitType.VAL)

focoos/cli/cli.py

@@ -73,7 +73,7 @@
 """
 
 import uuid
-from typing import Optional, cast, get_args
+from typing import Optional, Tuple, Union, cast, get_args
 
 import typer
 from typing_extensions import Annotated
@@ -100,6 +100,40 @@
 
 logger = get_logger("CLI")
 
+
+def parse_im_size(value: Union[str, int]) -> Union[int, Tuple[int, int]]:
+    """Parse image size from string or int.
+
+    Supports formats:
+    - int: 640 (square image)
+    - str: "640" (square image)
+    - str: "640,480" or "640x480" (non-square image as height,width)
+
+    Args:
+        value: Image size as int or string
+
+    Returns:
+        int for square images, tuple (height, width) for non-square
+    """
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        # Try comma or x separator
+        if "," in value:
+            parts = value.split(",")
+        elif "x" in value or "X" in value:
+            parts = value.replace("X", "x").split("x")
+        else:
+            # Single number - square image
+            return int(value)
+
+        if len(parts) == 2:
+            return (int(parts[0].strip()), int(parts[1].strip()))
+        else:
+            raise ValueError(f"Invalid image size format: {value}. Use '640', '640,480', or '640x480'")
+    return value
+
+
 app = typer.Typer(
     name="focoos",
     help=__doc__,
@@ -251,7 +285,9 @@ def train(
         Optional[str], typer.Option(help="Datasets directory (default: ~/FocoosAI/datasets/)")
     ] = None,
     dataset_layout: Annotated[DatasetLayout, typer.Option(help="Dataset layout")] = DatasetLayout.ROBOFLOW_COCO,
-    im_size: Annotated[int, typer.Option(help="Image size")] = 640,
+    im_size: Annotated[
+        str, typer.Option(help="Image size (int for square, or 'height,width' or 'heightxwidth' for non-square)")
+    ] = "640",
     output_dir: Annotated[Optional[str], typer.Option(help="Output directory")] = None,
     ckpt_dir: Annotated[Optional[str], typer.Option(help="Checkpoint directory")] = None,
     init_checkpoint: Annotated[Optional[str], typer.Option(help="Initial checkpoint path")] = None,
@@ -316,7 +352,9 @@ def train(
             generates a unique name using model name and UUID.
         datasets_dir (Optional[str]): Custom directory for datasets.
         dataset_layout (DatasetLayout): Layout format of the dataset. Defaults to ROBOFLOW_COCO.
-        im_size (int): Input image size for training. Defaults to 640.
+        im_size (str): Input image size for training. Can be int (e.g., "640") for square images,
+            or "height,width" or "heightxwidth" (e.g., "640,480" or "640x480") for non-square images.
+            Defaults to "640".
         output_dir (Optional[str]): Directory to save training outputs and logs.
         ckpt_dir (Optional[str]): Directory to save model checkpoints.
         init_checkpoint (Optional[str]): Path to initial checkpoint for transfer learning.
@@ -435,11 +473,12 @@ def train(
         validated_optimizer = cast(OptimizerType, optimizer.upper())
         assert optimizer in get_args(OptimizerType)
 
+        parsed_im_size = parse_im_size(im_size)
         train_command(
             model_name=model,
             dataset_name=dataset,
             dataset_layout=dataset_layout,
-            im_size=im_size,
+            im_size=parsed_im_size,
             run_name=run_name or f"{model}-{uuid.uuid4()}",
             output_dir=output_dir,
             ckpt_dir=ckpt_dir,
@@ -494,7 +533,9 @@ def val(
     ] = None,
     run_name: Annotated[Optional[str], typer.Option(help="Run name")] = None,
     dataset_layout: Annotated[DatasetLayout, typer.Option(help="Dataset layout")] = DatasetLayout.ROBOFLOW_COCO,
-    im_size: Annotated[int, typer.Option(help="Image size")] = 640,
+    im_size: Annotated[
+        str, typer.Option(help="Image size (int for square, or 'height,width' or 'heightxwidth' for non-square)")
+    ] = "640",
     output_dir: Annotated[Optional[str], typer.Option(help="Output directory")] = None,
     ckpt_dir: Annotated[Optional[str], typer.Option(help="Checkpoint directory")] = None,
     init_checkpoint: Annotated[Optional[str], typer.Option(help="Initial checkpoint")] = None,
@@ -677,11 +718,12 @@ def val(
         validated_optimizer = cast(OptimizerType, optimizer.upper())
         assert optimizer in get_args(OptimizerType)
 
+        parsed_im_size = parse_im_size(im_size)
         val_command(
             model_name=model,
             dataset_name=dataset,
             dataset_layout=dataset_layout,
-            im_size=im_size,
+            im_size=parsed_im_size,
             run_name=run_name or f"{model}-{uuid.uuid4()}",
             output_dir=output_dir,
             ckpt_dir=ckpt_dir,
@@ -881,7 +923,10 @@ def export(
     output_dir: Annotated[Optional[str], typer.Option(help="Output directory")] = None,
     device: Annotated[Optional[str], typer.Option(help="Device (cuda or cpu)")] = "cuda",
     onnx_opset: Annotated[Optional[int], typer.Option(help="ONNX opset version")] = 17,
-    im_size: Annotated[Optional[int], typer.Option(help="Image size for export")] = 640,
+    im_size: Annotated[
+        Optional[str],
+        typer.Option(help="Image size for export (int for square, or 'height,width' or 'heightxwidth' for non-square)"),
+    ] = "640",
     overwrite: Annotated[Optional[bool], typer.Option(help="Overwrite existing files")] = False,
 ):
     """Export a trained model to various deployment formats.
@@ -985,13 +1030,14 @@ def export(
     try:
         validated_device = cast(DeviceType, device)
         assert device in get_args(DeviceType)
+        parsed_im_size = parse_im_size(im_size) if im_size is not None else None
         export_command(
             model_name=model,
             format=format,
             output_dir=output_dir,
             device=validated_device,
             onnx_opset=onnx_opset,
-            im_size=im_size,
+            im_size=parsed_im_size,
             overwrite=overwrite,
         )
     except Exception as e:

focoos/cli/commands/export.py

@@ -55,7 +55,7 @@
     - [`focoos.ports.RuntimeType`][focoos.ports.RuntimeType]: Runtime type configurations
 """
 
-from typing import Literal, Optional
+from typing import Literal, Optional, Tuple, Union
 
 from focoos.model_manager import ModelManager
 from focoos.ports import ExportFormat, RuntimeType
@@ -70,7 +70,7 @@ def export_command(
     output_dir: Optional[str] = None,
     device: Optional[Literal["cuda", "cpu"]] = None,
     onnx_opset: Optional[int] = None,
-    im_size: Optional[int] = None,
+    im_size: Optional[Union[int, Tuple[int, int]]] = None,
     overwrite: Optional[bool] = None,
 ):
     """Export a model to different deployment formats.
@@ -114,7 +114,8 @@ def export_command(
         onnx_opset (Optional[int], optional): ONNX opset version for ONNX exports.
             Higher versions support more operations but may have compatibility issues.
             Common versions: 11, 13, 16, 17. Defaults to 17 if None.
-        im_size (Optional[int], optional): Input image size for the exported model.
+        im_size (Optional[Union[int, Tuple[int, int]]], optional): Input image size for the exported model.
+            If int, treated as square (size, size). If tuple, treated as (height, width).
             Used to define fixed input shapes for optimization.
             If None, uses the model's default input size.
             Defaults to None.

focoos/cli/commands/train.py

@@ -74,7 +74,7 @@
     - [`focoos.ports.TrainerArgs`][focoos.ports.TrainerArgs]: Training configuration
 """
 
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 from focoos.data.auto_dataset import AutoDataset
 from focoos.data.default_aug import get_default_by_task
@@ -101,7 +101,7 @@ def train_command(
     ## Dataset args
     dataset_name: str,
     dataset_layout: DatasetLayout,
-    im_size: int,
+    im_size: Union[int, Tuple[int, int]],
     ##################
     ## Training args
     run_name: str,
@@ -183,8 +183,9 @@ def train_command(
             and dataset identifiers.
         dataset_layout (DatasetLayout): Layout format of the dataset.
             Supported formats: ROBOFLOW_COCO, YOLO, COCO, etc.
-        im_size (int): Input image size for training. Images are resized
-            to this size while maintaining aspect ratio.
+        im_size (Union[int, Tuple[int, int]]): Input image size for training.
+            If int, treated as square (size, size). If tuple, treated as (height, width).
+            Images are resized to this size while maintaining aspect ratio.
         run_name (str): Unique name for this training run. Used for
             experiment tracking, logging, and output organization.
         output_dir (Optional[str], optional): Directory to save training outputs
@@ -342,8 +343,8 @@ def train_command(
         # Get augmentations
         logger.info("🎨 Setting up data augmentations")
         train_augs, val_augs = get_default_by_task(model.task, resolution=im_size or model.model_info.im_size)
-        train_dataset = auto_dataset.get_split(augs=train_augs.get_augmentations(), split=DatasetSplitType.TRAIN)
-        valid_dataset = auto_dataset.get_split(augs=val_augs.get_augmentations(), split=DatasetSplitType.VAL)
+        train_dataset = auto_dataset.get_split(augs=train_augs, split=DatasetSplitType.TRAIN)
+        valid_dataset = auto_dataset.get_split(augs=val_augs, split=DatasetSplitType.VAL)
 
         # Configure training arguments
         logger.info("⚙️ Configuring training parameters")

focoos/cli/commands/val.py

@@ -56,7 +56,7 @@
     - [`focoos.ports.TrainerArgs`][focoos.ports.TrainerArgs]: Configuration parameters
 """
 
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 from focoos.data.auto_dataset import AutoDataset
 from focoos.data.default_aug import get_default_by_task
@@ -82,7 +82,7 @@ def val_command(
     ## Dataset args
     dataset_name: str,
     dataset_layout: DatasetLayout,
-    im_size: int,
+    im_size: Union[int, Tuple[int, int]],
     ##################
     ## Training args
     run_name: str,
@@ -163,8 +163,9 @@ def val_command(
             and dataset identifiers.
         dataset_layout (DatasetLayout): Layout format of the dataset.
             Supported formats: ROBOFLOW_COCO, YOLO, COCO, etc.
-        im_size (int): Input image size for validation. Images are resized
-            to this size while maintaining aspect ratio.
+        im_size (Union[int, Tuple[int, int]]): Input image size for validation.
+            If int, treated as square (size, size). If tuple, treated as (height, width).
+            Images are resized to this size while maintaining aspect ratio.
         run_name (str): Unique name for this validation run. Used for
             result organization, logging, and report generation.
         output_dir (Optional[str], optional): Directory to save validation outputs
@@ -317,7 +318,7 @@ def val_command(
         # Get validation augmentations
         logger.info("🎨 Setting up validation augmentations")
         _, val_augs = get_default_by_task(task=model.model_info.task, resolution=im_size or model.model_info.im_size)
-        valid_dataset = auto_dataset.get_split(augs=val_augs.get_augmentations(), split=DatasetSplitType.VAL)
+        valid_dataset = auto_dataset.get_split(augs=val_augs, split=DatasetSplitType.VAL)
 
         # Configure validation arguments
         logger.info("⚙️ Configuring validation parameters")