diff --git a/backend/ai_server.py b/backend/ai_server.py index cb952ee..66cd1b5 100644 --- a/backend/ai_server.py +++ b/backend/ai_server.py @@ -183,9 +183,13 @@ _LABEL_ERROR = "" _DEVICE = os.environ.get("YOLO_DEVICE", "") _CONF = float(os.environ.get("YOLO_CONF", "0.25")) _POSE_CONF = float(os.environ.get("YOLO_POSE_CONF", "0.30")) -_POSE_RELIABLE_MIN_SCORE = float(os.environ.get("YOLO_POSE_RELIABLE_MIN_SCORE", "0.45")) +_POSE_RELIABLE_MIN_SCORE = float(os.environ.get("YOLO_POSE_RELIABLE_MIN_SCORE", "0.30")) _POSE_RELIABLE_MIN_KEYPOINTS = int(os.environ.get("YOLO_POSE_RELIABLE_MIN_KEYPOINTS", "6")) _POSE_RELIABLE_MIN_QUALITY = float(os.environ.get("YOLO_POSE_RELIABLE_MIN_QUALITY", "0.45")) +_POSITION_CONTEXT_MIN_SCORE = float(os.environ.get("YOLO_POSITION_CONTEXT_MIN_SCORE", "0.22")) +_POSITION_CONTEXT_MAX_SCORE = float(os.environ.get("YOLO_POSITION_CONTEXT_MAX_SCORE", "0.44")) +_POSITION_CONTEXT_BOOST_WEIGHT = float(os.environ.get("YOLO_POSITION_CONTEXT_BOOST_WEIGHT", "0.60")) +_POSITION_CONTEXT_OVERRIDE_MARGIN = float(os.environ.get("YOLO_POSITION_CONTEXT_OVERRIDE_MARGIN", "0.16")) _BATCH = int(os.environ.get("YOLO_BATCH", "16")) _IMGSZ = int(os.environ.get("YOLO_IMGSZ", "640")) _HALF = os.environ.get("YOLO_HALF", "0").lower() in {"1", "true", "yes", "on"} @@ -982,6 +986,79 @@ def combine_position_score(scores: dict[str, float], label: str, score: float) - scores[label] = clamp01(1 - (1 - current) * (1 - score)) +def append_prediction_source(prediction: dict, source: str) -> None: + source = str(source or "").strip() + if not source: + return + + current = str(prediction.get("source") or "").strip() + if not current: + prediction["source"] = source + return + + if source in {part.strip() for part in current.split("+")}: + return + + prediction["source"] = f"{current}+{source}" + + +def fuse_hybrid_position_scores( + pose_scores: dict[str, float], + context_scores: dict[str, float], +) -> tuple[str, float, bool, bool]: + labels = { + label + for label in set(pose_scores.keys()) | set(context_scores.keys()) + if not is_no_sex_position_label(label) + } + + best_position = "" + best_score = 0.0 + best_has_pose = False + best_has_context = False + best_pose_position = "" + best_pose_score = 0.0 + best_pose_has_context = False + + for label in labels: + pose_score = clamp01(pose_scores.get(label, 0.0)) + context_score = clamp01(context_scores.get(label, 0.0)) + has_pose = pose_score > 0 + has_context = context_score > 0 + + score = 0.0 + if has_pose: + score = pose_score + if has_context: + boost = clamp01(context_score * _POSITION_CONTEXT_BOOST_WEIGHT) + score = clamp01(1 - (1 - score) * (1 - boost)) + elif context_score >= _POSITION_CONTEXT_MIN_SCORE: + score = min(_POSITION_CONTEXT_MAX_SCORE, context_score) + + if score > best_score: + best_position = label + best_score = score + best_has_pose = has_pose + best_has_context = has_context + + if has_pose and score > best_pose_score: + best_pose_position = label + best_pose_score = score + best_pose_has_context = has_context + + if ( + best_pose_position + and not best_has_pose + and best_score <= best_pose_score + _POSITION_CONTEXT_OVERRIDE_MARGIN + ): + best_position = best_pose_position + best_score = best_pose_score + best_has_pose = True + best_has_context = best_pose_has_context + + return best_position, clamp01(best_score), best_has_pose, best_has_context + + def add_pose_pair_geometry_scores(scores: dict[str, float], persons: list[dict]) -> None: def add(label: str, score: float) -> None: combine_position_score(scores, label, score) @@ -1183,9 +1260,9 @@ def add_pose_scene_context_scores(scores: dict[str, float], prediction: dict, pe add("prone_bone", 0.07) -def best_pose_scene_position(prediction: dict, persons: list[dict]) -> tuple[str, float]: - scores: dict[str, float] = {} - has_direct_pose_position = False +def best_pose_scene_position(prediction: dict, persons: list[dict]) -> tuple[str, float, bool, bool]: + pose_scores: dict[str, float] = {} + context_scores: dict[str, float] = {} reliable_persons = reliable_pose_persons(persons) for person in reliable_persons: @@ -1195,43 +1272,42 @@ def best_pose_scene_position(prediction: dict, persons: list[dict]) -> tuple[str if is_no_sex_position_label(label) or label not in POSITION_LABELS: continue - has_direct_pose_position = True - combine_position_score(scores, label, score) + combine_position_score(pose_scores, label, score) quality = pose_keypoint_quality(person) if quality > 0: - combine_position_score(scores, label, 0.04 * quality) + combine_position_score(pose_scores, label, 0.04 * quality) - add_pose_scene_context_scores(scores, prediction, reliable_persons) + add_pose_scene_context_scores(context_scores, prediction, reliable_persons) - best_position = "" - best_score = 0.0 - for label, score in scores.items(): - if score > best_score: - best_position = label - best_score = score + best_position, best_score, has_pose_signal, has_context_signal = fuse_hybrid_position_scores( + pose_scores, + context_scores, + ) - if best_position and (has_direct_pose_position or best_score >= 0.30): - return best_position, clamp01(best_score) + if best_position and (has_pose_signal or best_score >= _POSITION_CONTEXT_MIN_SCORE): + return best_position, clamp01(best_score), has_pose_signal, has_context_signal - return "", 0.0 + return "", 0.0, False, has_context_signal def apply_pose_result_to_prediction(prediction: dict, result) -> dict: persons = pose_persons_from_result(result) - if not persons: - return prediction + if persons: + prediction["persons"] = persons - prediction["persons"] = persons - - best_position, best_score_value = best_pose_scene_position(prediction, persons) + best_position, best_score_value, _has_pose_signal, has_context_signal = best_pose_scene_position( + prediction, + persons, + ) if best_position: prediction["sexPosition"] = best_position prediction["sexPositionScore"] = best_score_value - source = str(prediction.get("source") or "").strip() - prediction["source"] = f"{source}+yolo_pose" if source else "yolo_pose" + append_prediction_source(prediction, "yolo_pose") + if has_context_signal: + append_prediction_source(prediction, "box_context") return prediction @@ -1241,6 +1317,19 @@ def apply_pose_batch_to_predictions(paths: list[str], predictions: list[dict], i current_pose_model = get_pose_model() if current_pose_model is None: + for prediction in predictions: + best_position, best_score_value, _has_pose_signal, has_context_signal = best_pose_scene_position( + prediction, + [], + ) + + if best_position: + prediction["sexPosition"] = best_position + prediction["sexPositionScore"] = best_score_value + + if has_context_signal: + append_prediction_source(prediction, "box_context") + return try: @@ -1255,6 +1344,19 @@ def apply_pose_batch_to_predictions(paths: list[str], predictions: list[dict], i ) except Exception as exc: _POSE_MODEL_ERROR = str(exc) + for prediction in predictions: + best_position, best_score_value, _has_pose_signal, has_context_signal = best_pose_scene_position( + prediction, + [], + ) + + if best_position: + prediction["sexPosition"] = best_position + prediction["sexPositionScore"] = best_score_value + + if has_context_signal: + append_prediction_source(prediction, "box_context") + return for prediction, pose_result in zip(predictions, pose_results): diff --git a/backend/analyze.go b/backend/analyze.go index 2e8519e..d391ccc 100644 --- a/backend/analyze.go +++ b/backend/analyze.go @@ -92,6 +92,22 @@ const ( analyzeMinComboScore = 0.36 ) +const ( + analyzePositionClipWindowSeconds = 3.0 + analyzePositionClipMinScore = 0.22 + analyzePositionClipMinFrames = 2 +) + +type analyzePositionEvidence struct { + Time float64 + Label string + Score float64 + Source string + PersonCount int + HasPose bool + HasContext bool +} + func analyzeVideoFrameFilter(intervalSeconds int) string { if intervalSeconds <= 0 { intervalSeconds = 1 @@ -1621,6 +1637,233 @@ func appendHighlightHitsFromPrediction( return append(hits, next...) } +func analyzePositionEvidenceFromPrediction( + pred TrainingPrediction, + t float64, +) (analyzePositionEvidence, bool) { + if !predictionUsableForAnalyze(pred) { + return analyzePositionEvidence{}, false + } + + label := strings.ToLower(strings.TrimSpace(pred.SexPosition)) + if isNoSexPositionLabel(label) || !isKnownPositionLabel(label) { + return analyzePositionEvidence{}, false + } + + score := pred.SexPositionScore + if score <= 0 { + score = 0.35 + } + score = clamp01(score) + if score < 0.16 { + return analyzePositionEvidence{}, false + } + + source := strings.ToLower(strings.TrimSpace(pred.Source)) + personCount := len(pred.Persons) + if personCount == 0 { + for _, box := range pred.Boxes { + if trainingIsPersonLikeLabel(box.Label) { + personCount++ + } + } + } + + return analyzePositionEvidence{ + Time: t, + Label: label, + Score: score, + Source: source, + PersonCount: personCount, + HasPose: strings.Contains(source, "yolo_pose") || len(pred.Persons) > 0, + HasContext: strings.Contains(source, "box_context") || len(pred.Boxes) > 0, + }, true +} + +func analyzePositionEvidenceWeight(item analyzePositionEvidence) float64 { + weight := 1.0 + + if item.HasPose && item.HasContext { + weight = 1.15 + } else if item.HasPose { + weight = 1.0 + } else if item.HasContext { + weight = 0.72 + } + + if item.PersonCount >= 2 { + weight += 0.08 + } + + return weight +} + +func buildClipPositionHitsFromEvidence( + evidence []analyzePositionEvidence, + duration float64, +) []analyzeHit { + if len(evidence) == 0 { + return []analyzeHit{} + } + + sort.SliceStable(evidence, func(i, j int) bool { + if evidence[i].Time != evidence[j].Time { + return evidence[i].Time < evidence[j].Time + } + return evidence[i].Label < evidence[j].Label + }) + + requiredFrames := analyzePositionClipMinFrames + if len(evidence) < requiredFrames { + requiredFrames = len(evidence) + } + if requiredFrames < 1 { + requiredFrames = 1 + } + + halfWindow := analyzePositionClipWindowSeconds / 2 + if halfWindow <= 0 { + halfWindow = math.Max(1, float64(analyzeVideoFrameIntervalSeconds)) + } + + type aggregate struct { + Label string + WeightedSum float64 + WeightSum float64 + Count int + PoseCount int + ContextCount int + Start float64 + End float64 + Marker float64 + BestScore float64 + } + + hits := make([]analyzeHit, 0, len(evidence)) + + for _, center := range evidence { + windowStart := center.Time - halfWindow + windowEnd := center.Time + halfWindow + if windowStart < 0 { + windowStart = 0 + } + if duration > 0 && windowEnd > duration { + windowEnd = duration + } + if windowEnd <= windowStart { + windowEnd = windowStart + math.Max(1, float64(analyzeVideoFrameIntervalSeconds)) + if duration > 0 && windowEnd > duration { + windowEnd = duration + } + } + + byLabel := map[string]*aggregate{} + + for _, item := range evidence { + if item.Time < windowStart-0.001 || item.Time > windowEnd+0.001 { + continue + } + + label := strings.ToLower(strings.TrimSpace(item.Label)) + if isNoSexPositionLabel(label) || !isKnownPositionLabel(label) { + continue + } + + score := clamp01(item.Score) + if score <= 0 { + continue + } + + agg := byLabel[label] + if agg == nil { + agg = &aggregate{ + Label: label, + Start: item.Time, + End: item.Time, + Marker: item.Time, + } + byLabel[label] = agg + } + + weight := analyzePositionEvidenceWeight(item) + agg.WeightedSum += score * weight + agg.WeightSum += weight + agg.Count++ + if item.HasPose { + agg.PoseCount++ + } + if item.HasContext { + agg.ContextCount++ + } + if item.Time < agg.Start { + agg.Start = item.Time + } + if item.Time > agg.End { + agg.End = item.Time + } + if score > agg.BestScore { + agg.BestScore = score + agg.Marker = item.Time + } + } + + var best *aggregate + bestScore := 0.0 + + for _, agg := range byLabel { + if agg.WeightSum <= 0 || agg.Count < requiredFrames { + continue + } + + avg := clamp01(agg.WeightedSum / agg.WeightSum) + stability := clamp01(float64(agg.Count) / math.Max(float64(requiredFrames), 3)) + sourceBonus := 0.0 + if agg.PoseCount > 0 && agg.ContextCount > 0 { + sourceBonus = 0.04 + } else if agg.PoseCount > 0 { + sourceBonus = 0.02 + } + + score := clamp01(avg*(0.86+0.14*stability) + sourceBonus) + if agg.PoseCount == 0 { + score = math.Min(score, trainingPositionContextMaxScore) + } + + if score < analyzePositionClipMinScore { + continue + } + + if score > bestScore { + best = agg + bestScore = score + } + } + + if best == nil { + continue + } + + start := math.Max(0, best.Start-halfWindow) + end := best.End + halfWindow + if duration > 0 { + end = math.Min(duration, end) + } + if end <= start { + end = math.Min(duration, start+math.Max(1, float64(analyzeVideoFrameIntervalSeconds))) + } + + hits = append(hits, analyzeHit{ + Time: best.Marker, + Label: "position:" + best.Label, + Score: bestScore, + Start: start, + End: end, + }) + } + + return mergeAnalyzeHits(hits) +} + func analyzeVideoFromFrames(ctx context.Context, outPath string) ([]analyzeHit, int64, int, error) { return analyzeVideoFromFramesForGoal(ctx, outPath) } @@ -1934,6 +2177,7 @@ func analyzeVideoFromFramesForGoal( batchOK := true detectorOnly := false + positionEvidence := []analyzePositionEvidence{} for startIdx := 0; startIdx < len(samples); startIdx += analyzeFramePredictBatchSize { if ctx.Err() != nil { @@ -2014,6 +2258,9 @@ func analyzeVideoFromFramesForGoal( } highlightHits = appendHighlightHitsFromPrediction(highlightHits, pred, sample.Time) + if item, ok := analyzePositionEvidenceFromPrediction(pred, sample.Time); ok { + positionEvidence = append(positionEvidence, item) + } } globalPercent := 50 + int(math.Round((float64(endIdx)/float64(total))*50)) @@ -2045,6 +2292,11 @@ func analyzeVideoFromFramesForGoal( ) } + highlightHits = append( + highlightHits, + buildClipPositionHitsFromEvidence(positionEvidence, durationSec)..., + ) + cleanHighlightHits := mergeAnalyzeHits(highlightHits) return cleanHighlightHits, startedAtMs, total, nil @@ -2063,6 +2315,9 @@ func analyzeVideoFromFramesForGoal( } highlightHits = appendHighlightHitsFromPrediction(highlightHits, pred, sample.Time) + if item, ok := analyzePositionEvidenceFromPrediction(pred, sample.Time); ok { + positionEvidence = append(positionEvidence, item) + } current := i + 1 @@ -2094,6 +2349,11 @@ func analyzeVideoFromFramesForGoal( ) } + highlightHits = append( + highlightHits, + buildClipPositionHitsFromEvidence(positionEvidence, durationSec)..., + ) + cleanHighlightHits := mergeAnalyzeHits(highlightHits) return cleanHighlightHits, startedAtMs, total, nil diff --git a/backend/training_test.go b/backend/training_test.go index 22cb73c..bd129af 100644 --- a/backend/training_test.go +++ b/backend/training_test.go @@ -393,6 +393,52 @@ func TestTrainingApplyPoseToPredictionBoostsPoseWithBoxContext(t *testing.T) { } } +func TestBuildClipPositionHitsFromEvidenceRequiresTemporalSupport(t *testing.T) { + hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{ + {Time: 1, Label: "doggy", Score: 0.55, HasPose: true, PersonCount: 2}, + {Time: 2, Label: "cowgirl", Score: 0.60, HasPose: true, PersonCount: 2}, + {Time: 6, Label: "doggy", Score: 0.58, HasPose: true, PersonCount: 2}, + }, 10) + + if len(hits) != 0 { + t.Fatalf("hits = %+v, want no temporally supported position", hits) + } +} + +func TestBuildClipPositionHitsFromEvidenceBuildsStablePosition(t *testing.T) { + hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{ + {Time: 1, Label: "doggy", Score: 0.45, HasPose: true, HasContext: true, PersonCount: 2}, + {Time: 2, Label: "doggy", Score: 0.47, HasPose: true, HasContext: true, PersonCount: 2}, + {Time: 3, Label: "doggy", Score: 0.44, HasPose: true, HasContext: true, PersonCount: 2}, + {Time: 6, Label: "cowgirl", Score: 0.62, HasPose: true, PersonCount: 2}, + }, 10) + + if len(hits) == 0 { + t.Fatal("expected stable doggy position hit") + } + if hits[0].Label != "position:doggy" { + t.Fatalf("label = %q, want position:doggy", hits[0].Label) + } + if hits[0].End <= hits[0].Start { + t.Fatalf("invalid hit span: %+v", hits[0]) + } +} + +func TestBuildClipPositionHitsFromEvidenceCapsContextOnlyScore(t *testing.T) { + hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{ + {Time: 1, Label: "missionary", Score: 0.90, HasContext: true, PersonCount: 2}, + {Time: 2, Label: "missionary", Score: 0.88, HasContext: true, PersonCount: 2}, + {Time: 3, Label: "missionary", Score: 0.92, HasContext: true, PersonCount: 2}, + }, 8) + + if len(hits) == 0 { + t.Fatal("expected context-only position hit") + } + if hits[0].Score > trainingPositionContextMaxScore { + t.Fatalf("score = %.3f, want <= %.3f", hits[0].Score, trainingPositionContextMaxScore) + } +} + func TestTrainingFilterPosePersonsByContextDropsUnannotatedPeople(t *testing.T) { persons := []TrainingPosePerson{ {