Add calibration for int8 quantization

yolo_nas_pose_to_onnx.py

@@ -12,6 +12,12 @@ import onnxruntime
 import os
 from super_gradients.training.utils.visualization.pose_estimation import PoseVisualization
 import matplotlib.pyplot as plt
+from datasets import load_dataset
+from torchvision import transforms
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
 
 os.environ['CRASH_HANDLER']='0'
 
@@ -19,7 +25,7 @@ os.environ['CRASH_HANDLER']='0'
 
 CONVERSION = True
 input_image_shape = [640, 640]
-quantization_modes = [None, ExportQuantizationMode.INT8, ExportQuantizationMode.FP16]
+quantization_modes = [ExportQuantizationMode.INT8, ExportQuantizationMode.FP16, None]
 output_predictions_format=DetectionOutputFormatMode.FLAT_FORMAT
 
 # NMS-related Setting
@@ -37,6 +43,61 @@ image_name = "https://deci-pretrained-models.s3.amazonaws.com/sample_images/beat
 # Check
 SHAPE_CHECK=True
 VISUAL_CHECK=True
+CALIBRATION_DATASET_CHECK=False
+
+# Function to convert tensor to image for visualization
+def tensor_to_image(tensor):
+    # Convert the tensor to a numpy array
+    numpy_image = tensor.numpy()
+
+    # The output of ToTensor() is in C x H x W format, convert to H x W x C
+    numpy_image = numpy_image.transpose(1, 2, 0)
+
+    # Undo the normalization (if any)
+    # numpy_image = numpy_image * std + mean  # Adjust based on your normalization
+
+    return numpy_image
+
+class HFDatasetWrapper(Dataset):
+    def __init__(self, hf_dataset, transform=None):
+        self.hf_dataset = hf_dataset
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.hf_dataset)
+
+    def __getitem__(self, idx):
+        item = self.hf_dataset[idx]
+        if self.transform:
+            item = self.transform(item)
+        return item['image']
+
+def preprocess(data):
+    # Convert byte data to PIL Image
+    image = data['image']
+
+    # Convert to RGB if not already
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+
+    # Define your transformations
+    transform = transforms.Compose([
+        transforms.Resize((640, 640)),  # Resize (example size)
+        transforms.ToTensor(),          # Convert to tensor
+        # Add normalization or other transformations if needed
+    ])
+
+	# Process Image
+    transformed = transform(image)
+
+    if CALIBRATION_DATASET_CHECK:
+		# Display the Processed Image
+        plt_image = tensor_to_image(transformed)
+        plt.imshow(plt_image)
+        plt.axis('off')  # Turn off axis numbers
+        plt.show()
+
+    return {'image': transformed}
 
 def iterate_over_flat_predictions(predictions, batch_size):
     [flat_predictions] = predictions
@@ -65,6 +126,11 @@ image = load_image(image_name)
 image = cv2.resize(image, (input_image_shape[1], input_image_shape[0]))
 image_bchw = np.transpose(np.expand_dims(image, 0), (0, 3, 1, 2))
 
+# Prepare Calibration Dataset for INT8 Quantization
+dataset = load_dataset("cppe-5", split="test")
+hf_dataset_wrapper = HFDatasetWrapper(dataset, transform=preprocess)
+calibration_loader = DataLoader(hf_dataset_wrapper, batch_size=8)
+
 for model_name in [Models.YOLO_NAS_POSE_L,  Models.YOLO_NAS_POSE_M,  Models.YOLO_NAS_POSE_N,  Models.YOLO_NAS_POSE_S ]:
 	for q in quantization_modes:
 		
@@ -94,7 +160,7 @@ for model_name in [Models.YOLO_NAS_POSE_L,  Models.YOLO_NAS_POSE_M,  Models.YOLO
 				engine=ExportTargetBackend.ONNXRUNTIME,
 				quantization_mode=q,
 				#selective_quantizer: Optional["SelectiveQuantizer"] = None,  # noqa
-				#calibration_loader: Optional[DataLoader] = None,
+				calibration_loader = calibration_loader,
 				#calibration_method: str = "percentile",
 				#calibration_batches: int = 16,
 				#calibration_percentile: float = 99.99,