Merge pull request #154 from FocoosAI/feat/export-quant

feat: add classification, quantization and simplify

Author: CuriousDolphin

docs/inference.md

@@ -1,4 +1,4 @@
-# How to Use a Computer Vision Model with Focoos
+# Inference with Focoos Models
 Focoos provides a powerful inference framework that makes it easy to deploy and use state-of-the-art computer vision models in production. Whether you're working on object detection, image classification, or other vision tasks, Focoos offers flexible deployment options that adapt to your specific needs.
 
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/FocoosAI/focoos/blob/main/tutorials/inference.ipynb)

docs/models/bisenetformer.md

@@ -121,18 +121,21 @@ Currently, you can find 3 bisenetformer models on the Focoos Hub, all for the se
 ### Quick Start with Pre-trained Model
 
 ```python
-from focoos.model_manager import ModelManager
+from focoos import ASSETS_DIR, ModelManager
+from PIL import Image
 
 # Load a pre-trained BisenetFormer model
 model = ModelManager.get("bisenetformer-m-ade")
 
 # Run inference on an image
-image = Image.open("path/to/image.jpg")
-result = model(image)
+image = ASSETS_DIR / "ADE_val_00000034"
+result = model.infer(image, threshold=0.5, annotate=True)
 
 # Process results
 for detection in result.detections:
     print(f"Class: {detection.label}, Confidence: {detection.conf:.3f}")
+# Visualize image
+Image.fromarray(result.image)
 ```
 
 ### Custom Model Configuration

docs/models/fai_cls.md

@@ -2,10 +2,25 @@
 
 ## Overview
 
-FAI-CLS is a versatile image classification model developed by FocoosAI that can utilize any backbone architecture for feature extraction. This model is designed for both single-label and multi-label image classification tasks, offering flexibility in architecture choices and training configurations.
+Fai-cls is a versatile image classification model developed by FocoosAI that can utilize any backbone architecture for feature extraction. This model is designed for both single-label and multi-label image classification tasks, offering flexibility in architecture choices and training configurations.
 
 The model employs a simple yet effective approach: a configurable backbone extracts features from input images, followed by a classification head that produces class predictions. This design enables easy adaptation to different domains and datasets while maintaining high performance and computational efficiency.
 
+## Available Models
+
+Currently, you can find 3 fai-cls models on the Focoos Hub, all trained on COCO dataset for image classification.
+
+| Model Name | Architecture | Domain (Classes) | Dataset | Metric | FPS Nvidia-T4 |
+|------------|--------------|------------------|----------|---------|--------------|
+| fai-cls-n-coco | Classification (STDC-Small) | Common Objects (80) | [COCO](https://cocodataset.org/#home) | F1: 48.66<br>Precision: 58.48<br>Recall: 41.66 | - |
+| fai-cls-s-coco | Classification (STDC-Small) | Common Objects (80) | [COCO](https://cocodataset.org/#home) | F1: 61.92<br>Precision: 68.69<br>Recall: 56.37 | - |
+| fai-cls-m-coco | Classification (STDC-Large) | Common Objects (80) | [COCO](https://cocodataset.org/#home) | F1: 66.98<br>Precision: 73.00<br>Recall: 61.88 | - |
+
+## Supported dataset
+- [ROBOFLOW_COCO](/focoos/api/ports/#focoos.ports.DatasetLayout) (multi-class)
+
+- [CLASSIFICATION_FOLDER](/focoos/api/ports/#focoos.ports.DatasetLayout)
+
 ## Neural Network Architecture
 
 The FAI-CLS architecture consists of two main components:
@@ -20,12 +35,14 @@ The FAI-CLS architecture consists of two main components:
 ### Classification Head
 - **Architecture**: Multi-layer perceptron (MLP) with configurable depth
 - **Components**:
+
   - Global Average Pooling (AdaptiveAvgPool2d) for spatial dimension reduction
   - Flatten layer to convert 2D features to 1D
   - Linear layers with ReLU activation
   - Dropout for regularization
   - Final linear layer for class predictions
 - **Configurations**:
+
   - **Single Layer**: Direct mapping from features to classes
   - **Two Layer**: Hidden layer with ReLU and dropout for better feature transformation
 
@@ -53,22 +70,26 @@ The FAI-CLS architecture consists of two main components:
 ### Single-Label Classification
 - **Output**: Single class prediction per image
 - **Use Cases**:
-  - Image categorization (animals, objects, scenes)
-  - Medical image diagnosis
-  - Quality control in manufacturing
-  - Content moderation
-  - Agricultural crop classification
+
+    - Image categorization (animals, objects, scenes)
+    - Medical image diagnosis
+    - Quality control in manufacturing
+    - Content moderation
+    - Agricultural crop classification
+
 - **Loss**: Cross-entropy or focal loss
 - **Configuration**: Set `multi_label=False`
 
 ### Multi-Label Classification
 - **Output**: Multiple class predictions per image
 - **Use Cases**:
-  - Multi-object recognition
-  - Image tagging and annotation
-  - Scene attribute recognition
-  - Medical condition classification
-  - Content-based image retrieval
+
+    - Multi-object recognition
+    - Image tagging and annotation
+    - Scene attribute recognition
+    - Medical condition classification
+    - Content-based image retrieval
+
 - **Loss**: Binary cross-entropy with logits
 - **Configuration**: Set `multi_label=True`
 
@@ -96,12 +117,6 @@ The model supports multiple loss function configurations:
 - **Features**: Optional label smoothing for better generalization
 - **Activation**: Softmax for probability distribution
 
-### Focal Loss
-- **Use Case**: Imbalanced datasets with hard-to-classify examples
-- **Parameters**:
-  - Alpha (α): Controls importance of rare class
-  - Gamma (γ): Focuses learning on hard examples
-- **Benefits**: Improved performance on imbalanced datasets
 
 ### Binary Cross-Entropy Loss
 - **Use Case**: Multi-label classification tasks
@@ -144,7 +159,30 @@ AdaptiveAvgPool2d(1) → Flatten → Linear(features → hidden_dim) → ReLU 
 This flexible architecture makes FAI-CLS suitable for a wide range of image classification applications, from simple binary classification to complex multi-label scenarios, while maintaining computational efficiency and ease of use.
 
 
-## Example Usage
+### Quick Start with Pre-trained Model
+
+```python
+from focoos import ASSETS_DIR, ModelManager
+from PIL import Image
+
+# Load a pre-trained model
+model = ModelManager.get("fai-cls-m-coco")
+
+image = ASSETS_DIR / "federer.jpg"
+result = model.infer(image,threshold=0.5, annotate=True)
+
+# Process results
+for detection in result.detections:
+    print(f"Class: {detection.label}, Confidence: {detection.conf:.3f}")
+
+# Visualize image
+Image.fromarray(result.image)
+
+```
+For the training process, please refer to the specific section of the documentation.
+
+
+## Custom Model Configuration
 
 ### Single-Label Classification Setup
 

docs/models/fai_detr.md

@@ -133,18 +133,22 @@ Currently, you can find 3 fai-detr models on the Focoos Hub, 2 trained on COCO a
 ### Quick Start with Pre-trained Model
 
 ```python
-from focoos.model_manager import ModelManager
+from focoos import ASSETS_DIR, ModelManager
+from PIL import Image
 
 # Load a pre-trained model
 model = ModelManager.get("fai-detr-m-coco")
 
 # Run inference on an image
-image = Image.open("path/to/image.jpg")
-result = model(image)
+image = ASSETS_DIR / "federer.jpg"
+result = model.infer(image,threshold=0.5, annotate=True)
 
 # Process results
 for detection in result.detections:
     print(f"Class: {detection.label}, Confidence: {detection.conf:.3f}")
+
+# Visualize image
+Image.fromarray(result.image)
 ```
 
 ### Custom Model Configuration

docs/models/fai_mf.md

@@ -115,18 +115,22 @@ Currently, you can find 5 fai-mf models on the Focoos Hub, 2 for semantic segmen
 ### Quick Start with Pre-trained Model
 
 ```python
-from focoos.model_manager import ModelManager
+from focoos import ASSETS_DIR, ModelManager
+from PIL import Image
 
 # Load a pre-trained BisenetFormer model
 model = ModelManager.get("fai-mf-l-ade")
 
 # Run inference on an image
-image = Image.open("path/to/image.jpg")
-result = model(image)
+image = ASSETS_DIR / "ADE_val_00000034"
+result = model.infer(image,threshold=0.5, annotate=True)
 
 # Process results
 for detection in result.detections:
     print(f"Class: {detection.label}, Confidence: {detection.conf:.3f}")
+
+# Visualize image
+Image.fromarray(result.image)
 ```
 
 ### Custom Model Configuration

docs/models/rtmo.md

@@ -127,18 +127,21 @@ The following RTMO models are available on the Focoos Hub for multi-person pose
 ```python
 from PIL import Image
 
-from focoos.model_manager import ModelManager
+from focoos import ModelManager, ASSETS_DIR
 
 # Load a pre-trained RTMO model
 model = ModelManager.get("rtmo-s-coco")
 
 # Run inference on an image
-image = Image.open("path/to/image.jpg")
-result = model.infer(image)
+image = ASSETS_DIR / "federer.jpg"
+result = model.infer(image,threshold=0.5, annotate=True)
 
 # Process results
 for detection in result.detections:
     print(f"Class: {detection.label}, Confidence: {detection.conf:.3f}")
+
+# Visualize image
+Image.fromarray(result.image)
 ```
 
 ### Custom Model Configuration

docs/quantization.md

@@ -0,0 +1,61 @@
+# Quantization (Beta)
+
+The quantization of Focoos models is currently in working in progress stage.
+
+currently tested and working for **classification models**.
+
+## Example
+
+```python
+from focoos import ModelManager,ASSETS_DIR, MODELS_DIR, RuntimeType
+from focoos.infer.quantizer import OnnxQuantizer, QuantizationCfg
+import os
+
+image_size = 224  # 224px input size
+model_name = "fai-cls-m-coco" # you can also take model from focoos hub with "hub://YOUR_MODEL_REF"
+im = ASSETS_DIR / "federer.jpg"
+
+model = ModelManager.get(model_name)
+
+exported_model = model.export(runtime_type=RuntimeType.ONNX_CPU, # optimized for edge or cpu
+    image_size=image_size,
+    dynamic_axes=False,  # quantization need static axes!
+    simplify_onnx=False, # simplify and optimize onnx model graph
+    onnx_opset=18,
+    out_dir=os.path.join(MODELS_DIR, "my_edge_model")) # save to models dir
+
+# benchmark onnx model
+exported_model.benchmark(iterations=100)
+
+# test onnx model
+
+result = exported_model.infer(im,annotate=True)
+Image.fromarray(result.image)
+
+
+quantization_cfg = QuantizationCfg(
+  size = image_size, # input size: must be same as exported model
+  calibration_images_folder = str(ASSETS_DIR), # Calibration images folder: It is strongly recommended
+                                               # to use the dataset validation split on which the model was trained.
+                                               # Here, for example, we will use the assets folder.
+  format="QDQ", # QO (QOperator): All the quantized operators have their own ONNX definitions, like QLinearConv, MatMulInteger etc.
+                # QDQ (Quantize-DeQuantize): inserts DeQuantizeLinear(QuantizeLinear(tensor)) between the original operators to simulate the quantization and dequantization process.
+  per_channel=True,      # Per-channel quantization: each channel has its own scale/zero-point → more accurate,
+                         # especially for convolutions, at the cost of extra memory and computation.
+  normalize_images=True, # normalize images during preprocessing: some models have normalization outside of model forward
+)
+
+quantizer = OnnxQuantizer(
+    input_model_path=exported_model.model_path,
+    cfg=quantization_cfg
+)
+model_path = quantizer.quantize(
+  benchmark=True # benchmark bot fp32 and int8 models
+)
+
+quantized_model = InferModel(model_path, runtime_type=RuntimeType.ONNX_CPU)
+
+res = quantized_model.infer(im,annotate=True)
+Image.fromarray(res.image)
+
+```

focoos/__init__.py

@@ -66,7 +66,6 @@
 from .utils.system import get_cpu_name, get_cuda_version, get_device_name, get_system_info
 from .utils.timer import took
 from .utils.vision import (
-    annotate_frame,
     annotate_image,
     base64mask_to_mask,
     binary_mask_to_base64,
@@ -165,7 +164,6 @@
     "get_device_name",
     "get_gpus_count",
     "get_cuda_version",
-    "annotate_frame",
     "annotate_image",
     "ModelRegistry",
     "InferLatency",

focoos/data/auto_dataset.py

@@ -22,7 +22,7 @@
     is_inside_sagemaker,
 )
 
-logger = get_logger(__name__)
+logger = get_logger("AutoDataset")
 
 
 class AutoDataset:
@@ -61,23 +61,23 @@ def __init__(
         self.dataset_path = str(dataset_path)
         self.dataset_name = dataset_name
         logger.info(
-            f"✅ Dataset name: {self.dataset_name}, Dataset Path: {self.dataset_path}, Dataset Layout: {self.layout}"
+            f"🔄 Loading dataset {self.dataset_name}, 📁 Dataset Path: {self.dataset_path}, 🗂️ Dataset Layout: {self.layout}"
         )
 
     def _load_split(self, dataset_name: str, split: DatasetSplitType) -> DictDataset:
         if self.layout == DatasetLayout.CATALOG:
-            return DictDataset.from_catalog(ds_name=dataset_name, split=split, root=self.dataset_path)
+            return DictDataset.from_catalog(ds_name=dataset_name, split_type=split, root=self.dataset_path)
         else:
             ds_root = self.dataset_path
             if not check_folder_exists(ds_root):
                 raise FileNotFoundError(f"Dataset {ds_root} not found")
             split_path = self._get_split_path(dataset_root=ds_root, split_type=split)
             if self.layout == DatasetLayout.ROBOFLOW_SEG:
-                return DictDataset.from_roboflow_seg(ds_dir=split_path, task=self.task)
+                return DictDataset.from_roboflow_seg(ds_dir=split_path, task=self.task, split_type=split)
             elif self.layout == DatasetLayout.CLS_FOLDER:
-                return DictDataset.from_folder(root_dir=split_path)
+                return DictDataset.from_folder(root_dir=split_path, split_type=split)
             elif self.layout == DatasetLayout.ROBOFLOW_COCO:
-                return DictDataset.from_roboflow_coco(ds_dir=split_path, task=self.task)
+                return DictDataset.from_roboflow_coco(ds_dir=split_path, task=self.task, split_type=split)
             else:  # Focoos
                 raise NotImplementedError(f"Dataset layout {self.layout} not implemented")
 

focoos/data/converters.py

@@ -17,7 +17,7 @@
 
 from focoos.data.datasets.dict_dataset import DictDataset
 from focoos.data.transforms.resize_short_length import resize_shortest_length
-from focoos.ports import DatasetMetadata, Task
+from focoos.ports import DatasetMetadata, DatasetSplitType, Task
 from focoos.utils.logger import get_logger
 from focoos.utils.system import list_files_with_extensions
 
@@ -450,10 +450,14 @@ def convert_datasetninja_to_mask_dataset(
         )
 
     task = Task.SEMSEG
-    train_dataset = DictDataset.from_segmentation(ds_dir=os.path.join(dataset_path, train_split_name), task=task)
+    train_dataset = DictDataset.from_segmentation(
+        ds_dir=os.path.join(dataset_path, train_split_name), task=task, split_type=DatasetSplitType.TRAIN
+    )
     logger.info(f"Train dataset: {train_dataset}")
 
-    val_dataset = DictDataset.from_segmentation(ds_dir=os.path.join(dataset_path, val_split_name), task=task)
+    val_dataset = DictDataset.from_segmentation(
+        ds_dir=os.path.join(dataset_path, val_split_name), task=task, split_type=DatasetSplitType.VAL
+    )
     logger.info(f"Val dataset: {val_dataset}")
 
     for split in [(train_dataset, "train"), (val_dataset, "val")]: