Added new dataset

scikit-learn/convert2onnx.py

@@ -4,18 +4,27 @@
 import argparse
 import joblib
 
-from sklearn.datasets import fetch_california_housing, load_diabetes, load_iris
 from skl2onnx import convert_sklearn
 from skl2onnx.common.data_types import FloatTensorType
 
 
 def load_dataset(dataset_name):
     if dataset_name == 'california':
+        from sklearn.datasets import fetch_california_housing
         dataset = fetch_california_housing()
     elif dataset_name == 'diabetes':
+        from sklearn.datasets import load_diabetes
         dataset = load_diabetes()
     elif dataset_name == 'iris':
+        from sklearn.datasets import load_iris
         dataset = load_iris()
+    elif dataset_name == "cardiotocography":
+        from sklearn.datasets import fetch_openml
+        dataset = fetch_openml(name=dataset_name, version=1, as_frame=False)
+        X, y = dataset.data, dataset.target
+        s = y == "3"
+        y = s.astype(int)
+        return X, y
     else:
         raise ValueError("Invalid dataset name")
     return dataset.data, dataset.target
@@ -53,7 +62,7 @@ python convert2onnx.py california adaboost_regressor.joblib adaboost_regressor.o
 """
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Converts a sklearn model to ONNX format.')
-    parser.add_argument('dataset_name', type=str, help='Name of the dataset. Choose from: "california", "diabetes", or "iris".')
+    parser.add_argument('dataset_name', type=str, help='Name of the dataset. Choose from: "california", "diabetes", "iris" or "cardiotocography".')
     parser.add_argument('model_path', type=str, help='Path to the trained model file.')
     parser.add_argument('onnx_filename', type=str, help='The filename for the output ONNX file.')
     args = parser.parse_args()