File size: 3,537 Bytes
4bab068
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
TriageResult: the single structured object that flows out of the pipeline.
Keeping this as one dataclass makes app.py, tests, and the report generator
all agree on one contract instead of passing loose dicts around.
"""

from dataclasses import dataclass, field
from typing import List, Dict, Any
from .config import risk_level_from_score


@dataclass
class TriageResult:
    scene_summary: List[str] = field(default_factory=list)
    location_type: str = "unknown"
    water_level: str = "n/a"
    road_status: str = "unknown"
    structural_damage: str = "unknown"
    hazards: List[str] = field(default_factory=list)
    people_visible_count: int = 0
    people_description: str = ""
    risk_score: int = 0
    risk_level: str = ""
    recommended_actions: List[str] = field(default_factory=list)
    confidence: float = 0.0
    yolo_detections: Dict[str, int] = field(default_factory=dict)
    raw_model_output: str = ""

    def __post_init__(self):
        # Clamp risk score defensively — never trust a free-text model blindly.
        self.risk_score = max(0, min(100, int(self.risk_score)))
        if not self.risk_level:
            self.risk_level, _ = risk_level_from_score(self.risk_score)

    @property
    def risk_color(self) -> str:
        _, color = risk_level_from_score(self.risk_score)
        return color

    @classmethod
    def from_dict(cls, data: Dict[str, Any], yolo_detections: Dict[str, int],
                   raw_model_output: str = "") -> "TriageResult":
        """Build a TriageResult from the VLM's (possibly messy) JSON dict,
        filling safe defaults for any missing/malformed field."""
        people = data.get("people_status", {}) or {}
        return cls(
            scene_summary=list(data.get("scene_summary", []) or []),
            location_type=str(data.get("location_type", "unknown")),
            water_level=str(data.get("water_level", "n/a")),
            road_status=str(data.get("road_status", "unknown")),
            structural_damage=str(data.get("structural_damage", "unknown")),
            hazards=list(data.get("hazards", []) or []),
            people_visible_count=int(people.get("visible_count", 0) or 0),
            people_description=str(people.get("description", "")),
            risk_score=int(data.get("risk_score", 0) or 0),
            risk_level=str(data.get("risk_level", "") or ""),
            recommended_actions=list(data.get("recommended_actions", []) or []),
            confidence=float(data.get("confidence", 0.0) or 0.0),
            yolo_detections=yolo_detections,
            raw_model_output=raw_model_output,
        )

    def to_markdown(self) -> str:
        bullets = "\n".join(f"- {s}" for s in self.scene_summary) or "- (no details extracted)"
        hazards = "\n".join(f"- ⚠️ {h}" for h in self.hazards) or "- None detected"
        actions = "\n".join(f"- [ ] {a}" for a in self.recommended_actions) or "- None"
        detections = ", ".join(f"{v}x {k}" for k, v in self.yolo_detections.items()) or "none"

        return f"""## Scene Summary
{bullets}

**Location type:** {self.location_type}  **Water level:** {self.water_level}
**Road status:** {self.road_status}  **Structural damage:** {self.structural_damage}
**People visible:** {self.people_visible_count}{self.people_description or "n/a"}

### Hazards
{hazards}

### Risk Score: {self.risk_score}/100  ({self.risk_level})
*Model confidence: {self.confidence:.0%}*

### Recommended Response
{actions}

---
**YOLO grounded detections:** {detections}
"""