Spaces:
Running
Running
File size: 4,576 Bytes
188709e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import os
import logging
from typing import Iterator, Dict, Any, Optional
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("WaveUIDataset")
try:
import fiftyone as fo
from fiftyone.utils.huggingface import load_from_hub
except ImportError:
logger.warning("FiftyOne or HuggingFace module is not installed. To run this dataset loader, install them using: pip install fiftyone")
class WaveUIDatasetManager:
"""
Manager to load, stream, and parse Voxel51's WaveUI-25K dataset.
Used to train Aura Assist's UI component visual understanding, layout detection, and visual agent reasoning.
Dataset Reference: "Voxel51/WaveUI-25k"
"""
def __init__(self, dataset_name: str = "WaveUI-25k"):
self.dataset_name = dataset_name
def load_waveui_dataset(self, max_samples: Optional[int] = None) -> Any:
"""
Loads the Voxel51/WaveUI-25k dataset directly from the Hugging Face hub using FiftyOne.
Optionally limits the sample size using max_samples.
"""
logger.info(f"Loading Voxel51/WaveUI-25k dataset (Max Samples: {max_samples or 'All'})...")
try:
dataset = load_from_hub(
"Voxel51/WaveUI-25k",
max_samples=max_samples
)
logger.info("WaveUI-25k dataset loaded successfully into FiftyOne!")
return dataset
except Exception as e:
logger.error(f"Error loading Voxel51/WaveUI-25k dataset: {e}")
return None
def launch_explorer_session(self, max_samples: Optional[int] = 100) -> Any:
"""
Launches the interactive FiftyOne desktop app session to explore UI bounding boxes visually.
"""
dataset = self.load_waveui_dataset(max_samples=max_samples)
if dataset:
logger.info("Launching FiftyOne Interactive Dashboard App...")
session = fo.launch_app(dataset)
return session
return None
def _parse_sample_features(self, sample: Any) -> Dict[str, Any]:
"""
Parses a single FiftyOne sample into standard JSON format with coordinates, bounding boxes,
and labels suitable for Aura Assist element detection pipelines.
"""
filepath = sample.filepath
tags = sample.tags
elements = []
# Typically, FiftyOne object detection datasets store bounding boxes in standard fields like 'ground_truth' or 'detections'
detections_field = sample.get_field("ground_truth") or sample.get_field("detections")
if detections_field and hasattr(detections_field, "detections"):
for detection in detections_field.detections:
label = detection.label
bbox = detection.bounding_box # FiftyOne format: [top-left-x, top-left-y, width, height] (normalized 0-1)
elements.append({
"label": label,
"confidence": getattr(detection, "confidence", 1.0),
"box_normalized": {
"x": bbox[0],
"y": bbox[1],
"width": bbox[2],
"height": bbox[3]
}
})
return {
"image_path": filepath,
"tags": tags,
"elements": elements,
"metadata": {
"source": "Voxel51/WaveUI-25k",
"width": sample.metadata.width if sample.metadata else None,
"height": sample.metadata.height if sample.metadata else None
}
}
def stream_samples(self, max_samples: Optional[int] = None) -> Iterator[Dict[str, Any]]:
"""
Streams structured FiftyOne samples with layout component detections sequentially.
"""
dataset = self.load_waveui_dataset(max_samples=max_samples)
if not dataset:
logger.warning("No active dataset loaded. Aborting stream.")
return
logger.info("Streaming WaveUI-25k samples...")
for sample in dataset:
try:
yield self._parse_sample_features(sample)
except Exception as e:
logger.error(f"Skipping corrupt WaveUI-25k sample: {e}")
continue
# Quick test execution scaffold
if __name__ == "__main__":
manager = WaveUIDatasetManager()
logger.info("WaveUI-25k Dataset Manager Scaffolding initialized.")
print("WaveUI-25k Dataset Ready for visual component training!")
|