include enrichment

Dockerfile

@@ -10,7 +10,7 @@ RUN pip install pip -U
 
 RUN pip install renumics-spotlight==1.3.0rc8 pyarrow
 
-RUN pip install datasets
+RUN pip install datasets cleanvision
 
 COPY . .
 RUN mkdir -p /code/.cache

prepare.py

@@ -2,7 +2,18 @@ import pickle
 import datasets
 import os
 import pandas as pd
+from cleanvision.imagelab import Imagelab
 
+def cv_issues_cleanvision(df, image_name='image'):
+
+    image_paths = df['image'].to_list()
+    imagelab = Imagelab(filepaths=image_paths)
+    imagelab.find_issues()
+
+    df_cv=imagelab.issues.reset_index()
+
+    return df_cv
+    
 if __name__ == "__main__":
     cache_file = "dataset_cache.parquet"
     if os.path.exists(cache_file):
@@ -14,7 +25,12 @@ if __name__ == "__main__":
         dataset = datasets.load_dataset("renumics/cifar100-enriched", split="test")
         print("Dataset loaded using datasets.load_dataset().")
         
-        df = dataset.to_pandas()       
+        df = dataset.to_pandas()
+        
+        df=df.reset_index(drop=True)
+        
+        df_cv=cv_issues_cleanvision(df)
+        df = pd.concat([df, df_cv], axis=1)    
        
 
         # Save dataset to cache