Update pipeline_full.py

pipeline_full.py

@@ -311,9 +311,7 @@ def step_siglip_clustering(video_path: str, out_dir: str) -> Dict[str, str]:
             imageElement.style.display = 'block';
             placeholderText.style.display = 'none';
         }}
-
         var chartElement = document.getElementById('scatter-plot-3d');
-
         chartElement.on('plotly_click', function(data) {{
             var customdata = data.points[0].customdata;
             displayImage(customdata);
@@ -741,36 +739,360 @@ def step_ball_path(video_path: str, out_dir: str) -> Dict[str, Any]:
     }
 
 
-# -------------------- 8. stats-only process_video --------------------
+# -------------------- 8. NEW: full-match analysis + event-annotated video --------------------
 
 
-def process_video_stats(video_path: str) -> Dict[str, Any]:
+def step_analyze_and_annotate_video(video_path: str, out_dir: str) -> Dict[str, Any]:
+    """
+    Single pass over the video that:
+      * tracks players & ball
+      * computes distance & speed per player (pitch coordinates)
+      * estimates ball possession per team & per player
+      * detects simple events (pass, tackle/interception, clearance, shot)
+      * renders an annotated MP4 with overlays
+    """
     ensure_models_loaded()
+    os.makedirs(out_dir, exist_ok=True)
+
+    video_info = sv.VideoInfo.from_video_path(video_path)
+    fps = video_info.fps
+    dt = 1.0 / fps
 
     tracker = sv.ByteTrack()
     tracker.reset()
+
+    # homography smoothing
+    Ms = deque(maxlen=5)
+
+    # stats
+    distance_covered_m = defaultdict(float)        # tid -> meters
+    possession_time_player = defaultdict(float)    # tid -> seconds
+    possession_time_team = defaultdict(float)      # team_id -> seconds
+    team_of_player = {}                            # tid -> team_id
+    events: List[Dict[str, Any]] = []
+
+    # last positions for speed / distance
+    last_pitch_pos: Dict[int, np.ndarray] = {}
+    prev_owner_tid: Optional[int] = None
+    prev_ball_pos_pitch: Optional[np.ndarray] = None
+
+    # simple goal centers in pitch coordinates (x is length, y is width)
+    goal_centers = {
+        0: np.array([0.0, PITCH_CONFIG.width / 2.0]),
+        1: np.array([PITCH_CONFIG.length, PITCH_CONFIG.width / 2.0]),
+    }
+
+    # annotators
+    ellipse_annotator = sv.EllipseAnnotator(
+        color=sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "#FFD700"]),
+        thickness=2,
+    )
+    label_annotator = sv.LabelAnnotator(
+        color=sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "#FFD700"]),
+        text_color=sv.Color.from_hex("#000000"),
+        text_position=sv.Position.BOTTOM_CENTER,
+    )
+    triangle_annotator = sv.TriangleAnnotator(
+        color=sv.Color.from_hex("#FFD700"), base=25, height=21, outline_thickness=1
+    )
+
+    sink_path = os.path.join(out_dir, "annotated_events.mp4")
+    sink = sv.VideoSink(sink_path, video_info)
+
+    # text overlay control
+    current_event_text = ""
+    event_text_frames_left = 0
+    EVENT_TEXT_DURATION_S = 2.0
+    EVENT_TEXT_DURATION_FRAMES = int(EVENT_TEXT_DURATION_S * fps)
+
+    frame_generator = sv.get_video_frames_generator(video_path)
+
+    with sink:
+        for frame_idx, frame in enumerate(tqdm(frame_generator, total=video_info.total_frames,
+                                               desc="analyze + annotate")):
+            t = frame_idx * dt
+
+            # --- detections + tracking ---
+            det_result = PLAYER_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
+            detections = sv.Detections.from_inference(det_result)
+
+            ball_dets = detections[detections.class_id == BALL_ID]
+            ball_dets.xyxy = sv.pad_boxes(xyxy=ball_dets.xyxy, px=10)
+
+            non_ball = detections[detections.class_id != BALL_ID]
+            non_ball = non_ball.with_nms(threshold=0.5, class_agnostic=True)
+            tracked = tracker.update_with_detections(non_ball)
+
+            goalkeepers_dets = tracked[tracked.class_id == GOALKEEPER_ID]
+            players_dets = tracked[tracked.class_id == PLAYER_ID]
+            referees_dets = tracked[tracked.class_id == REFEREE_ID]
+
+            # --- field homography ---
+            field_result = FIELD_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
+            key_points = sv.KeyPoints.from_inference(field_result)
+            filt = key_points.confidence[0] > 0.5
+            frame_ref = key_points.xy[0][filt]
+            pitch_ref = np.array(PITCH_CONFIG.vertices)[filt]
+
+            transformer = ViewTransformer(source=frame_ref, target=pitch_ref)
+            Ms.append(transformer.m)
+            transformer.m = np.mean(np.array(Ms), axis=0)
+
+            # --- team classification & pitch positions ---
+            frame_players_xy_pitch = None
+            frame_ball_pos_pitch = None
+
+            if len(players_dets) > 0:
+                crops = [sv.crop_image(frame, xyxy) for xyxy in players_dets.xyxy]
+                team_preds = TEAM_CLASSIFIER.predict(crops)
+                players_dets.class_id = team_preds  # now class_id = team_id (0/1)
+
+                frame_players_xy_img = players_dets.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+                frame_players_xy_pitch = transformer.transform_points(points=frame_players_xy_img)
+
+                for tid, team_id, pos_pitch in zip(
+                    players_dets.tracker_id, players_dets.class_id, frame_players_xy_pitch
+                ):
+                    tid_int = int(tid)
+                    team_of_player[tid_int] = int(team_id)
+
+                    prev_pos = last_pitch_pos.get(tid_int)
+                    speed_kmh = 0.0
+                    if prev_pos is not None:
+                        dist_m = float(np.linalg.norm(pos_pitch - prev_pos))
+                        distance_covered_m[tid_int] += dist_m
+                        speed_kmh = (dist_m / dt) * 3.6  # m/s -> km/h
+                    last_pitch_pos[tid_int] = pos_pitch
+
+            if len(ball_dets) > 0:
+                frame_ball_xy_img = ball_dets.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+                frame_ball_xy_pitch = transformer.transform_points(points=frame_ball_xy_img)
+                frame_ball_pos_pitch = frame_ball_xy_pitch[0]
+
+            # --- possession owner ---
+            owner_tid: Optional[int] = None
+            POSSESSION_RADIUS_M = 5.0
+
+            if frame_ball_pos_pitch is not None and frame_players_xy_pitch is not None:
+                dists = np.linalg.norm(frame_players_xy_pitch - frame_ball_pos_pitch, axis=1)
+                j = int(np.argmin(dists))
+                if dists[j] < POSSESSION_RADIUS_M:
+                    owner_tid = int(players_dets.tracker_id[j])
+
+            # accumulate possession time
+            if owner_tid is not None:
+                possession_time_player[owner_tid] += dt
+                owner_team = team_of_player.get(owner_tid)
+                if owner_team is not None:
+                    possession_time_team[owner_team] += dt
+
+            # --- event detection (simple heuristics) ---
+            def register_event(ev: Dict[str, Any], text: str):
+                nonlocal current_event_text, event_text_frames_left
+                events.append(ev)
+                current_event_text = text
+                event_text_frames_left = EVENT_TEXT_DURATION_FRAMES
+
+            # possession change events, passes, tackles, interceptions
+            if owner_tid != prev_owner_tid:
+                if owner_tid is not None and prev_owner_tid is not None:
+                    prev_team = team_of_player.get(prev_owner_tid)
+                    cur_team = team_of_player.get(owner_tid)
+
+                    travel_m = 0.0
+                    if prev_ball_pos_pitch is not None and frame_ball_pos_pitch is not None:
+                        travel_m = float(np.linalg.norm(frame_ball_pos_pitch - prev_ball_pos_pitch))
+
+                    MIN_PASS_TRAVEL_M = 3.0
+
+                    if prev_team is not None and cur_team is not None:
+                        if prev_team == cur_team and travel_m > MIN_PASS_TRAVEL_M:
+                            # pass
+                            register_event(
+                                {
+                                    "type": "pass",
+                                    "t": float(t),
+                                    "from_tid": int(prev_owner_tid),
+                                    "to_tid": int(owner_tid),
+                                    "team_id": int(cur_team),
+                                    "extra": {"distance_m": travel_m},
+                                },
+                                f"Pass: #{prev_owner_tid} → #{owner_tid} (Team {cur_team})",
+                            )
+                        elif prev_team != cur_team:
+                            # tackle vs interception
+                            d_pp = 999.0
+                            if frame_players_xy_pitch is not None:
+                                # get current positions
+                                pos_prev = last_pitch_pos.get(int(prev_owner_tid))
+                                pos_cur = last_pitch_pos.get(int(owner_tid))
+                                if pos_prev is not None and pos_cur is not None:
+                                    d_pp = float(np.linalg.norm(pos_prev - pos_cur))
+                            ev_type = "tackle" if d_pp < 3.0 else "interception"
+                            label = "Tackle" if ev_type == "tackle" else "Interception"
+                            register_event(
+                                {
+                                    "type": ev_type,
+                                    "t": float(t),
+                                    "from_tid": int(prev_owner_tid),
+                                    "to_tid": int(owner_tid),
+                                    "team_id": int(cur_team),
+                                    "extra": {"player_distance_m": d_pp, "ball_travel_m": travel_m},
+                                },
+                                f"{label}: #{owner_tid} wins ball from #{prev_owner_tid}",
+                            )
+
+                # generic possession-change event
+                register_event(
+                    {
+                        "type": "possession_change",
+                        "t": float(t),
+                        "from_tid": int(prev_owner_tid) if prev_owner_tid is not None else None,
+                        "to_tid": int(owner_tid) if owner_tid is not None else None,
+                        "team_id": int(team_of_player.get(owner_tid)) if owner_tid is not None else None,
+                        "extra": {},
+                    },
+                    "" if owner_tid is None else f"Team {team_of_player.get(owner_tid)} in possession",
+                )
+
+            # shot / clearance based on ball speed & direction
+            if (
+                prev_ball_pos_pitch is not None
+                and frame_ball_pos_pitch is not None
+                and owner_tid is not None
+            ):
+                v = (frame_ball_pos_pitch - prev_ball_pos_pitch) / dt  # m/s
+                speed_mps = float(np.linalg.norm(v))
+                speed_kmh = speed_mps * 3.6
+                HIGH_SPEED_KMH = 18.0
+
+                if speed_kmh > HIGH_SPEED_KMH:
+                    shooter_team = team_of_player.get(owner_tid)
+                    if shooter_team is not None:
+                        target_goal = goal_centers[1 - shooter_team]
+                        direction = target_goal - frame_ball_pos_pitch
+                        cos_angle = float(
+                            np.dot(v, direction)
+                            / (np.linalg.norm(v) * np.linalg.norm(direction) + 1e-6)
+                        )
+
+                        if cos_angle > 0.8:
+                            register_event(
+                                {
+                                    "type": "shot",
+                                    "t": float(t),
+                                    "from_tid": int(owner_tid),
+                                    "to_tid": None,
+                                    "team_id": int(shooter_team),
+                                    "extra": {"speed_kmh": speed_kmh},
+                                },
+                                f"Shot by #{owner_tid} (Team {shooter_team}) – {speed_kmh:.1f} km/h",
+                            )
+                        else:
+                            register_event(
+                                {
+                                    "type": "clearance",
+                                    "t": float(t),
+                                    "from_tid": int(owner_tid),
+                                    "to_tid": None,
+                                    "team_id": int(shooter_team),
+                                    "extra": {"speed_kmh": speed_kmh},
+                                },
+                                f"Clearance by #{owner_tid} (Team {shooter_team})",
+                            )
+
+            prev_owner_tid = owner_tid
+            prev_ball_pos_pitch = frame_ball_pos_pitch
+
+            # --- frame drawing ---
+            annotated = frame.copy()
+
+            # build labels for players: id + speed + distance
+            player_labels: List[str] = []
+            if frame_players_xy_pitch is not None and len(players_dets) > 0:
+                for tid, pos_pitch in zip(players_dets.tracker_id, frame_players_xy_pitch):
+                    tid_int = int(tid)
+                    prev_pos = last_pitch_pos.get(tid_int)
+                    speed_kmh = 0.0
+                    if prev_pos is not None:
+                        dist_m = float(np.linalg.norm(pos_pitch - prev_pos))
+                        speed_kmh = (dist_m / dt) * 3.6
+                    d_total = distance_covered_m[tid_int]
+                    team_id = team_of_player.get(tid_int, -1)
+                    player_labels.append(
+                        f"#{tid_int} T{team_id} {speed_kmh:4.1f} km/h {d_total:.1f} m"
+                    )
+
+                annotated = ellipse_annotator.annotate(
+                    scene=annotated, detections=players_dets
+                )
+                annotated = label_annotator.annotate(
+                    scene=annotated, detections=players_dets, labels=player_labels
+                )
+
+            # draw ball
+            annotated = triangle_annotator.annotate(scene=annotated, detections=ball_dets)
+
+            # --- HUD: possession percentages ---
+            total_poss_time = sum(possession_time_team.values()) + 1e-6
+            team0_pct = 100.0 * possession_time_team.get(0, 0.0) / total_poss_time
+            team1_pct = 100.0 * possession_time_team.get(1, 0.0) / total_poss_time
+
+            hud_text = f"Team 0 Ball Control: {team0_pct:5.2f}%   Team 1 Ball Control: {team1_pct:5.2f}%"
+            cv2.rectangle(
+                annotated,
+                (20, annotated.shape[0] - 60),
+                (annotated.shape[1] - 20, annotated.shape[0] - 20),
+                (255, 255, 255),
+                -1,
+            )
+            cv2.putText(
+                annotated,
+                hud_text,
+                (30, annotated.shape[0] - 30),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.8,
+                (0, 0, 0),
+                2,
+                cv2.LINE_AA,
+            )
+
+            # --- event banner ---
+            if event_text_frames_left > 0 and current_event_text:
+                cv2.rectangle(annotated, (20, 20), (annotated.shape[1] - 20, 90), (255, 255, 255), -1)
+                cv2.putText(
+                    annotated,
+                    current_event_text,
+                    (30, 70),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1.0,
+                    (0, 0, 0),
+                    2,
+                    cv2.LINE_AA,
+                )
+                event_text_frames_left -= 1
+
+            sink.write_frame(annotated)
+
+    # finalize stats
+    total_poss = sum(possession_time_team.values()) + 1e-6
+    possession_percent_team = {
+        int(team): 100.0 * t_sec / total_poss for team, t_sec in possession_time_team.items()
+    }
+
     stats = {
-        "distance_covered": defaultdict(float),
-        "team_classifications": {},
-        "field_key_points": None,
+        "distance_covered_m": {str(tid): float(d) for tid, d in distance_covered_m.items()},
+        "possession_time_player_s": {str(tid): float(t_sec) for tid, t_sec in possession_time_player.items()},
+        "possession_time_team_s": {int(team): float(t_sec) for team, t_sec in possession_time_team.items()},
+        "possession_percent_team": possession_percent_team,
+        "team_of_player": {str(tid): int(team) for tid, team in team_of_player.items()},
     }
 
-    frame_gen = sv.get_video_frames_generator(video_path)
-    for frame_idx, frame in enumerate(frame_gen):
-        result = PLAYER_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
-        detections = sv.Detections.from_inference(result)
-        field_result = FIELD_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
-        field_key_points = sv.KeyPoints.from_inference(field_result)
-        if stats["field_key_points"] is None:
-            stats["field_key_points"] = field_key_points.xy.tolist()
-        detections = tracker.update_with_detections(detections)
-        player_crops = [sv.crop_image(frame, xyxy) for xyxy in detections.xyxy]
-        if player_crops:
-            preds = TEAM_CLASSIFIER.predict(player_crops)
-            for tid, team in zip(detections.tracker_id, preds):
-                stats["team_classifications"][str(tid)] = int(team)
-    stats["distance_covered"] = dict(stats["distance_covered"])
-    return stats
+    return {
+        "annotated_video": sink_path,
+        "stats": stats,
+        "events": events,
+    }
 
 
 # -------------------- 9. full pipeline entrypoint --------------------
@@ -787,29 +1109,33 @@ def run_full_pipeline(video_path: str, job_dir: str) -> Dict[str, Any]:
 
     os.makedirs(job_dir, exist_ok=True)
 
-    update_progress("siglip", 0.15, "Running SigLIP clustering...")
+    update_progress("siglip", 0.10, "Running SigLIP clustering...")
     siglip_out = step_siglip_clustering(video_path, os.path.join(job_dir, "siglip"))
 
-    update_progress("team_classifier", 0.30, "Training TeamClassifier...")
+    update_progress("team_classifier", 0.25, "Training TeamClassifier...")
     train_team_classifier_on_video(video_path)
 
-    update_progress("basic_frames", 0.45, "Generating basic annotated frames...")
+    update_progress("basic_frames", 0.35, "Generating basic annotated frames...")
     basic_paths = step_basic_frames(video_path, os.path.join(job_dir, "frames"))
 
-    update_progress("advanced_views", 0.60, "Generating advanced radar / Voronoi views...")
+    update_progress("advanced_views", 0.45, "Generating advanced radar / Voronoi views...")
     adv_paths = step_single_frame_advanced(video_path, os.path.join(job_dir, "advanced"))
 
-    update_progress("ball_path", 0.80, "Computing ball path and heatmap...")
+    update_progress("ball_path", 0.60, "Computing ball path and heatmap...")
     ball_paths = step_ball_path(video_path, os.path.join(job_dir, "ball_path"))
 
-    update_progress("stats", 0.90, "Calculating stats...")
-    stats = process_video_stats(video_path)
+    update_progress("events_video", 0.80, "Analyzing match and rendering event-annotated video...")
+    analysis_out = step_analyze_and_annotate_video(
+        video_path, os.path.join(job_dir, "analysis")
+    )
 
     result = {
         "basic": basic_paths,
         "advanced": adv_paths,
         "ball": ball_paths,
-        "stats": stats,
+        "stats": analysis_out["stats"],
+        "events": analysis_out["events"],
+        "annotated_video": analysis_out["annotated_video"],
         "siglip_html": siglip_out["plot_html"],
     }