Fixed code to account for the hugging face dataset.

app.py

@@ -1,36 +1,32 @@
+import io
 import os
+import tarfile
 from pathlib import Path
 
 import gradio as gr
 import pandas as pd
+from datasets import load_dataset
 from fastai.vision.all import *
 
+# Load the dataset from Hugging Face
+dataset = load_dataset("freddyal/smb_collection")
 
-def create_dataframe(image_dir, csv_path):
-    df = pd.read_csv(csv_path)
-    df['file'] = df['file'].apply(lambda x: str(Path(image_dir) / x))
+df = pd.DataFrame(dataset['train'])
 
-    df['file_exists'] = df['file'].apply(os.path.isfile)
+def label_generation(file_name):
+    if "gen2_images" in file_name:
+        return "Gen 2"
+    elif "gen3_images" in file_name:
+        return "Gen 3"
+    else:
+        return "Non SMB"
 
-    missing_files = df[~df['file_exists']]
-    if not missing_files.empty:
-        print("Warning: The following files are missing:")
-        print(missing_files['file'])
 
-    df = df[df['file_exists']].drop('file_exists', axis=1)
+df['label'] = df['__key__'].apply(label_generation)
 
-    return df
+df_gen2 = df[df['label'] == 'Gen 2']
 
-gen2_images_csv = "./gen2_images.csv"
-gen2_images_dir = "./gen2_images"
-df_gen2 = create_dataframe(gen2_images_dir, gen2_images_csv)
-
-gen3_images_csv = "./gen3_images.csv"
-gen3_images_dir = "./gen3_images"
-df_gen3 = create_dataframe(gen3_images_dir, gen3_images_csv)
-
-df_gen2['label'] = 'Gen 2'
-df_gen3['label'] = 'Gen 3'
+df_gen3 = df[df['label'] == 'Gen 3']
 
 def create_neither_dataframe(path):
     file_paths = []
@@ -75,4 +71,4 @@ interface = gr.Interface(
     live=True
 )
 
-interface.launch(share=True)
+interface.launch(share=True)

requirements.txt

@@ -1,3 +1,4 @@
 gradio~=4.44.0
 fastai~=2.7.17
-pandas~=2.2.0
+pandas~=2.2.0
+datasets~=3.0.0