bdanko/loaf_resolution_512

from datasets import load_dataset
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

# 1. Load the dataset
# Note: Since this is a private repo, ensure you have run `huggingface-cli login`
repo_id = "bdanko/loaf_resolution_512"
print(f"Downloading {repo_id}...")
dataset = load_dataset(repo_id, split="train")

# 2. Grab the first example
example = dataset[0]

# Hugging Face automatically decodes the Parquet bytes into a PIL Image
img = example["image"] 
objects = example["objects"]

# Print image-level metadata
print("\n--- Image Metadata ---")
print(f"File Name: {example['file_name']}")
print(f"Image ID: {example['image_id']}")
print(f"Dimensions: {example['width']}x{example['height']}")
print(f"Camera Height (Image): {example['camera_height_img']}")

# 3. Set up the matplotlib plot
fig, ax = plt.subplots(1, figsize=(10, 8))
ax.imshow(img)
ax.axis("off")
ax.set_title(f"Sample: {example['file_name']}", fontsize=14)

print("\n--- Object Metadata ---")

# 4. Iterate through all objects in this image
# We use zip to unpack all the parallel lists stored inside the 'objects' dictionary
num_objects = len(objects["annotation_id"])

for i in range(num_objects):
    ann_id = objects["annotation_id"][i]
    cat_id = objects["category_id"][i]
    bbox = objects["bbox"][i]
    seg = objects["segmentation"][i]
    
    # Custom / 3D annotations
    ritbox = objects["ritbox"][i]
    rot_box = objects["rotated_box"][i]
    world_loc = objects["world_location"][i]
    person_loc = objects["person_location"][i]
    cam_height_ann = objects["camera_height_ann"][i]
    
    print(f"\nObject {i+1} (Ann ID: {ann_id} | Category: {cat_id})")
    print(f"  - 2D BBox: {bbox}")
    print(f"  - Ritbox: {ritbox}")
    print(f"  - Rotated Box: {rot_box}")
    print(f"  - World Location: {world_loc}")
    print(f"  - Person Location: {person_loc}")
    print(f"  - Camera Height: {cam_height_ann}")

    # --- Plotting the 2D Bounding Box ---
    # COCO bbox format is [top_left_x, top_left_y, width, height]
    if len(bbox) == 4:
        x, y, w, h = bbox
        rect = patches.Rectangle(
            (x, y), w, h, 
            linewidth=2, edgecolor='red', facecolor='none', label=f"Cat {cat_id}"
        )
        ax.add_patch(rect)
        ax.text(x, y - 5, f"Cat {cat_id}", color='red', fontsize=10, weight='bold')

    # --- Plotting the Segmentation Polygon ---
    # COCO segmentation is a flat list: [x1, y1, x2, y2, ...]
    if len(seg) > 0:
        # Reshape flat list into a list of (x, y) coordinate pairs
        poly_coords = np.array(seg).reshape(-1, 2)
        polygon = patches.Polygon(
            poly_coords, 
            linewidth=2, edgecolor='cyan', facecolor='cyan', alpha=0.3
        )
        ax.add_patch(polygon)

plt.tight_layout()
plt.show()