added lots datasets

app.py

@@ -37,31 +37,79 @@ from ncut_pytorch.backbone import MODEL_DICT, LAYER_DICT, RES_DICT
 from ncut_pytorch import NCUT
 from ncut_pytorch import eigenvector_to_rgb, rotate_rgb_cube
 
-DATASET_TUPS = [
-    # (name, num_classes)
-    ('UCSC-VLAA/Recap-COCO-30K', None),
-    ('nateraw/pascal-voc-2012', None),
-    ('johnowhitaker/imagenette2-320', 10),
-    ('jainr3/diffusiondb-pixelart', None),
-    ('nielsr/CelebA-faces', None),
-    ('JapanDegitalMaterial/Places_in_Japan', None),
-    ('Borismile/Anime-dataset', None),
-    ('Multimodal-Fatima/CUB_train', 200),
-    ('mrm8488/ImageNet1K-val', 1000),
-    ("trashsock/hands-images", 8),
-]
-DATASET_NAMES = [tup[0] for tup in DATASET_TUPS]
-DATASET_CLASSES = [tup[1] for tup in DATASET_TUPS]
+
+DATASETS = {
+    'Common': [
+        ('mrm8488/ImageNet1K-val', 1000),
+        ('UCSC-VLAA/Recap-COCO-30K', None),
+        ('nateraw/pascal-voc-2012', None),
+        ('johnowhitaker/imagenette2-320', 10),
+        ('Multimodal-Fatima/CUB_train', 200),
+        ('saragag/FlBirds', 7),
+        ('microsoft/cats_vs_dogs', 2),
+        ('Robotkid2696/food_classification', 20),
+    ],
+    'Face': [
+        ('nielsr/CelebA-faces', None),
+        ('huggan/anime-faces', None),
+    ],
+    'Pose': [
+        ('razdab/sign_pose_M', None),
+        ('sayakpaul/poses-controlnet-dataset', None),
+        ('EgoThink/EgoThink', None),
+        ('junjuice0/vtuber-tachi-e', None),
+        ('Fiacre/small-animal-poses-controlnet-dataset', None),
+    ],
+    'Hand': [
+        ('trashsock/hands-images', 8),
+        ('dduka/guitar-chords-v3', None),
+    ],
+    'Satellite': [
+        ('arakesh/deepglobe-2448x2448', None),
+        ('tanganke/eurosat', 10),
+        ('wangyi111/EuroSAT-SAR', None),
+        ('efoley/sar_tile_512', None),
+    ],
+    'Medical': [
+        ('Mahadih534/Chest_CT-Scan_images-Dataset', 4),
+        ('Falah/Alzheimer_MRI', 4),
+        ('sartajbhuvaji/Brain-Tumor-Classification', 4),
+        ('TrainingDataPro/chest-x-rays', None),
+        ('hongrui/mimic_chest_xray_v_1', None),
+        ('Leonardo6/path-vqa', None),
+        ('Itsunori/path-vqa_jap', None),
+        ('ruby-jrl/isic-2024-2', None),
+        ('VRJBro/lung_cancer_dataset', 5),
+        ('keremberke/blood-cell-object-detection', None)
+    ],
+    'Miscs': [
+        ('yashvoladoddi37/kanjienglish', None),
+        ('Borismile/Anime-dataset', None),
+        ('jainr3/diffusiondb-pixelart', None),
+        ('jlbaker361/dcgan-eval-creative_gan_256_256', None),
+        ('Francesco/csgo-videogame', None),
+        ('Francesco/apex-videogame', None),
+        ('Marqo/deepfashion-multimodal', None),
+        ('huggan/pokemon', None),
+        ('huggan/few-shot-universe', None),
+        ('huggan/flowers-102-categories', None),
+        ('huggan/inat_butterflies_top10k', None),
+    ]
+}
+CENTER_CROP_DATASETS = ["razdab/sign_pose_M"]
+
 
 from datasets import load_dataset
 
 def download_all_datasets():
-    for name in DATASET_NAMES:
-        print(f"Downloading {name}")
-        try:
-            load_dataset(name, trust_remote_code=True)
-        except Exception as e:
-            print(f"Error downloading {name}: {e}")
+    for cat in DATASETS.keys():
+        for tup in DATASETS[cat]:
+            name = tup[0]
+            print(f"Downloading {name}")
+            try:
+                load_dataset(name, trust_remote_code=True)
+            except Exception as e:
+                print(f"Error downloading {name}: {e}")
 
 def compute_ncut(
     features,
@@ -1126,21 +1174,31 @@ def make_dataset_images_section(advanced=False, is_random=False):
 
     gr.Markdown('### Load Datasets')
     load_images_button = gr.Button("🔴 Load Images", elem_id="load-images-button", variant='primary')
-    advanced_radio = gr.Radio(["Basic", "Advanced"], label="Datasets", value="Advanced" if advanced else "Basic", elem_id="advanced-radio")
+    advanced_radio = gr.Radio(["Basic", "Advanced"], label="Datasets", value="Advanced" if advanced else "Basic", elem_id="advanced-radio", show_label=True)
     with gr.Column() as basic_block:
-        example_gallery = gr.Gallery(value=example_items, label="Example Set A", show_label=False, columns=[3], rows=[2], object_fit="scale-down", height="200px", show_share_button=False, elem_id="example-gallery")
+        example_gallery = gr.Gallery(value=example_items, label="Example Images", show_label=True, columns=[3], rows=[2], object_fit="scale-down", height="200px", show_share_button=False, elem_id="example-gallery")
     with gr.Column() as advanced_block:
-        dataset_names = DATASET_NAMES
-        dataset_classes = DATASET_CLASSES
+        # dataset_names = DATASET_NAMES
+        # dataset_classes = DATASET_CLASSES
+        dataset_categories = list(DATASETS.keys())
+        defualt_cat = dataset_categories[0]
+        def get_choices(cat):
+            return [tup[0] for tup in DATASETS[cat]]
+        defualt_choices = get_choices(defualt_cat)
         with gr.Row():
-            dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="mrm8488/ImageNet1K-val", elem_id="dataset", min_width=300)
+            dataset_radio = gr.Radio(dataset_categories, label="Dataset Category", value=defualt_cat, elem_id="dataset-radio", show_label=True, min_width=600)
+            # dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="mrm8488/ImageNet1K-val", elem_id="dataset", min_width=300)
+            dataset_dropdown = gr.Dropdown(defualt_choices, label="Dataset name", value=defualt_choices[0], elem_id="dataset", min_width=400)
+            dataset_radio.change(fn=lambda x: gr.update(choices=get_choices(x), value=get_choices(x)[0]), inputs=dataset_radio, outputs=dataset_dropdown)
             # num_images_slider = gr.Number(10, label="Number of images", elem_id="num_images")
-            num_images_slider = gr.Slider(1, 1000, step=1, label="Number of images", value=10, elem_id="num_images")
+            num_images_slider = gr.Slider(1, 1000, step=1, label="Number of images", value=10, elem_id="num_images", min_width=200)
             if not is_random:
                 filter_by_class_checkbox = gr.Checkbox(label="Filter by class", value=True, elem_id="filter_by_class_checkbox")
                 filter_by_class_text = gr.Textbox(label="Class to select", value="0,33,99", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. (1000 classes)", visible=True)
-                is_random_checkbox = gr.Checkbox(label="Random shuffle", value=False, elem_id="random_seed_checkbox")
-                random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=False)
+                # is_random_checkbox = gr.Checkbox(label="Random shuffle", value=False, elem_id="random_seed_checkbox")
+                # random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=False)
+                is_random_checkbox = gr.Checkbox(label="Random shuffle", value=True, elem_id="random_seed_checkbox")
+                random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=True)
             if is_random:
                 filter_by_class_checkbox = gr.Checkbox(label="Filter by class", value=False, elem_id="filter_by_class_checkbox")
                 filter_by_class_text = gr.Textbox(label="Class to select", value="0,33,99", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. (1000 classes)", visible=False)
@@ -1159,10 +1217,17 @@ def make_dataset_images_section(advanced=False, is_random=False):
     advanced_radio.change(fn=lambda x: gr.update(visible=x=="Advanced"), inputs=advanced_radio, outputs=[advanced_block])
     advanced_radio.change(fn=lambda x: gr.update(visible=x=="Basic"), inputs=advanced_radio, outputs=[basic_block])
     
+    def find_num_classes(dataset_name):
+        num_classes = None
+        for cat, datasets in DATASETS.items():
+            datasets = [tup[0] for tup in datasets]
+            if dataset_name in datasets:
+                num_classes = DATASETS[cat][datasets.index(dataset_name)][1]
+                break
+        return num_classes
     
     def change_filter_options(dataset_name):
-        idx = dataset_names.index(dataset_name)
-        num_classes = dataset_classes[idx]
+        num_classes = find_num_classes(dataset_name)
         if num_classes is None:
             return (gr.Checkbox(label="Filter by class", value=False, elem_id="filter_by_class_checkbox", visible=False),
                     gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info="e.g. `0,1,2`. This dataset has no class label", visible=False))
@@ -1171,8 +1236,7 @@ def make_dataset_images_section(advanced=False, is_random=False):
     dataset_dropdown.change(fn=change_filter_options, inputs=dataset_dropdown, outputs=[filter_by_class_checkbox, filter_by_class_text])
     
     def change_filter_by_class(is_filter, dataset_name):
-        idx = dataset_names.index(dataset_name)
-        num_classes = dataset_classes[idx]
+        num_classes = find_num_classes(dataset_name)
         return gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. ({num_classes} classes)", visible=is_filter)
     filter_by_class_checkbox.change(fn=change_filter_by_class, inputs=[filter_by_class_checkbox, dataset_dropdown], outputs=filter_by_class_text)
     
@@ -1227,9 +1291,26 @@ def make_dataset_images_section(advanced=False, is_random=False):
                 image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
             else:
                 image_idx = list(range(num_images))
-        images = [dataset[i]['image'] for i in image_idx]
+        key = 'image' if 'image' in dataset[0] else list(dataset[0].keys())[0]
+        images = [dataset[i][key] for i in image_idx]
         gr.Info(f"Loaded {len(images)} images from {dataset_name}")
         del dataset
+        
+        if dataset_name in CENTER_CROP_DATASETS:
+            def center_crop_image(image):
+                # image: PIL image
+                w, h = img.size
+                min_hw = min(h, w)
+                # center crop
+                left = (w - min_hw) // 2
+                top = (h - min_hw) // 2
+                right = left + min_hw
+                bottom = top + min_hw
+                print(left, top, right, bottom)
+                img = img.crop((left, top, right, bottom))
+                return image
+            images = [center_crop_image(image) for image in images]
+        
         return images   
     
     load_images_button.click(load_dataset_images,