Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import tinygrad | |
| class Manager: | |
| def __init__(self) -> None: | |
| self._classes = [] | |
| self._class_to_idx = {} | |
| self._samples = [] | |
| print("Manager initialized") | |
| def handle_index_imagefolder(self, root: str = "/dataset") -> None: | |
| classes = sorted( | |
| d for d in os.listdir(root) | |
| if os.path.isdir(os.path.join(root, d)) | |
| ) | |
| class_to_idx = {c: i for i, c in enumerate(classes)} | |
| samples = [] | |
| for c in classes: | |
| cdir = os.path.join(root, c) | |
| for name in os.listdir(cdir): | |
| if name.lower().endswith((".jpg", ".jpeg", ".png")): | |
| samples.append((os.path.join(cdir, name), class_to_idx[c])) | |
| print(f"Indexed ImageFolder: {len(samples)} samples found in {len(classes)} classes.") | |
| self._classes = classes | |
| self._class_to_idx = class_to_idx | |
| self._samples = samples | |
| def handle_get_random_sample(self) -> dict: | |
| if not self._samples: | |
| raise RuntimeError("No samples indexed. Please upload an image folder first.") | |
| idx = np.random.randint(0, len(self._samples)) | |
| path, class_idx = self._samples[idx] | |
| class_name = self._classes[class_idx] | |
| return {"imagePath": path, "label": class_name} | |