Spaces:
Runtime error
Runtime error
File size: 5,422 Bytes
efaa228 | 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """
Annotation QA Environment — Type-Safe Models.
Defines the API contract for the Annotation QA Environment:
- AnnotationQAAction: What corrections the agent can make
- AnnotationQAObservation: What the agent sees (image + annotations)
- AnnotationQAState: Episode metadata
The agent reviews intentionally-flawed annotations on real COCO val2017 images
and must fix bounding boxes, correct class labels, add missing annotations,
or remove spurious ones. A VLM (Vision-Language Model) is used to visually
inspect the images.
"""
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field
# ──────────────────────────────────────────────
# Annotation data structure
# ──────────────────────────────────────────────
class Annotation(BaseModel):
"""A single annotation: bounding box + class label."""
id: int
bbox: List[float] = Field(
...,
description="Bounding box as [x, y, w, h] normalized to 0.0–1.0",
min_length=4,
max_length=4,
)
class_label: str = Field(..., description="Object class label, e.g. 'car', 'person'")
# ──────────────────────────────────────────────
# Action
# ──────────────────────────────────────────────
class AnnotationQAAction(BaseModel):
"""
An action the agent can take to correct annotations.
action_type determines which fields are required:
- "adjust_bbox": requires annotation_id, new_bbox
- "change_class": requires annotation_id, new_class
- "add_annotation": requires new_bbox, new_class
- "remove_annotation": requires annotation_id
- "submit": no extra fields needed (finalizes episode)
"""
action_type: Literal[
"adjust_bbox",
"change_class",
"remove_annotation",
"add_annotation",
"submit",
"flag_safety",
"change_attribute",
"flag_missing",
]
annotation_id: Optional[int] = Field(
None, description="ID of the annotation to modify"
)
new_bbox: Optional[List[float]] = Field(
None,
description="New bounding box [x, y, w, h] in 0.0–1.0",
min_length=4,
max_length=4,
)
new_class: Optional[str] = Field(
None, description="New class label"
)
new_attribute: Optional[str] = Field(
None, description="New attribute description for an object"
)
missing_class: Optional[str] = Field(
None, description="Class of an object that was missing bounding boxes"
)
metadata: Dict[str, Any] = Field(default_factory=dict)
# ──────────────────────────────────────────────
# Observation
# ──────────────────────────────────────────────
class AnnotationQAObservation(BaseModel):
"""
What the agent sees after each step.
Includes the image URL, scene description, current annotations (some may
be wrong), available classes, and progress info. The VLM agent uses the
image_url to visually inspect the scene.
"""
done: bool = False
reward: Optional[float] = None
# Image information (real COCO val2017)
image_url: Optional[str] = Field(
None, description="Public URL to the COCO val2017 image"
)
image_width: int = Field(0, description="Image width in pixels")
image_height: int = Field(0, description="Image height in pixels")
# Scene information
scene_description: str = Field(
"", description="Natural-language description of the scene and its objects"
)
scene_objects: List[Dict[str, Any]] = Field(
default_factory=list,
description="Ground-truth object list with positions (visible to agent as scene context)",
)
# Current annotations (may contain errors)
annotations: List[Annotation] = Field(
default_factory=list,
description="Current annotations the agent should review/fix",
)
# Task context
available_classes: List[str] = Field(
default_factory=list,
description="Valid class labels for this task (COCO 80 categories)",
)
task_id: str = ""
task_description: str = ""
# Progress
corrections_made: int = 0
step_count: int = 0
max_steps: int = 20
# Feedback
message: str = ""
last_action_error: Optional[str] = None
# ──────────────────────────────────────────────
# State
# ──────────────────────────────────────────────
class AnnotationQAState(BaseModel):
"""Episode metadata — internal state tracked by the environment."""
episode_id: Optional[str] = None
step_count: int = 0
task_id: str = ""
sample_id: str = ""
initial_quality: float = 0.0
current_quality: float = 0.0
corrections_made: int = 0
|