PhysiGo/app/lib/exercises/exercises_validation/widgets/pose_detector.dart

import 'dart:async';

import 'package:body_detection/models/pose_landmark.dart';
import 'package:physigo/exercises/exercises_validation/models/exercise.dart';
import 'package:rxdart/rxdart.dart';
import 'package:body_detection/body_detection.dart';
import 'package:body_detection/models/image_result.dart';
import 'package:body_detection/models/pose.dart';
import 'package:flutter/material.dart';

import 'exercise_indicator.dart';
import 'pose_painter.dart';

typedef MeanFilteredData = List<List<double>>;
typedef LandmarkVariations = List<List<double>>;

enum StepExercise {
  notInPlace,
  ready,
  start,
  end,
}

class PoseDetector extends StatefulWidget {
  final Exercise exercise;
  const PoseDetector({required this.exercise, Key? key}) : super(key: key);

  @override
  State<PoseDetector> createState() => _PoseDetectorState();
}

class _PoseDetectorState extends State<PoseDetector> {
  static const meanFilterBuffer = 5;
  final List<StreamSubscription> _streamSubscriptions = [];
  final List<StreamController> _streamControllers = [];
  final StreamController<Pose> _poseController = StreamController.broadcast();
  final StreamController<StepExercise> _stepExerciseController = StreamController.broadcast();
  late final Stream<List<PoseLandmark>> _exerciseJointsStream;
  Image? _cameraImage;
  Pose? _detectedPose;
  Size _imageSize = Size.zero;
  late Future<void> _startCamera;
  late Stream<MeanFilteredData> _meanFilterStream;
  late Stream<int> _repCounter;

  @override
  initState() {
    super.initState();
    _streamControllers.add(_poseController);
    _streamControllers.add(_stepExerciseController);
    _exerciseJointsStream = _getExerciseJointsStream(_poseController.stream);
    CombineLatestStream.combine2(
        _exerciseJointsStream
            .map((event) => event.where((e) => widget.exercise.jointsOnScreen.contains(e.type)).toList()),
        _stepExerciseController.stream,
        (a, b) => [a, b]).listen(_handleNotInPlacePosition);
    CombineLatestStream.combine2(
        _exerciseJointsStream
            .map((event) => event.where((e) => widget.exercise.jointsOnScreen.contains(e.type)).toList())
            .bufferCount(5, 1),
        _stepExerciseController.stream,
        (a, b) => [a, b]).listen(_handleReadyPosition);
    _startCamera = _startCameraStream();
    _meanFilterStream = _getMeanFilterStream(_exerciseJointsStream);
    _streamSubscriptions.add(
      CombineLatestStream.combine2(_stepExerciseController.stream, _meanFilterStream, (a, b) => [a, b]).listen(
        (value) {
          final stepExercise = value.first as StepExercise;
          if (stepExercise == StepExercise.notInPlace) return;
          final meanFilteredData = value.last as MeanFilteredData;
          final isStartOfExerciseMovement = widget.exercise.isAtStartMovement(meanFilteredData);
          final isEndOfExerciseMovement = widget.exercise.isAtEndMovement(meanFilteredData);
          if ((stepExercise == StepExercise.start) && isEndOfExerciseMovement) {
            _stepExerciseController.add(StepExercise.end);
          }
          if ((stepExercise == StepExercise.end || stepExercise == StepExercise.ready) && isStartOfExerciseMovement) {
            _stepExerciseController.add(StepExercise.start);
          }
        },
      ),
    );
    _stepExerciseController.add(StepExercise.notInPlace);
    _repCounter = _stepExerciseController.stream
        .pairwise()
        .where((event) => event.first == StepExercise.end && event.last == StepExercise.start)
        .scan((int accumulated, value, index) => accumulated + 1, 0);
  }

  void _handleNotInPlacePosition(List<Object?> event) {
    final poseLandmarks = event.first as List<PoseLandmark>;
    final stepExercise = event.last as StepExercise;
    if (stepExercise == StepExercise.notInPlace) return;
    for (final poseLandmark in poseLandmarks) {
      if (poseLandmark.inFrameLikelihood < 0.8) {
        _stepExerciseController.add(StepExercise.notInPlace);
        return;
      }
    }
  }

  void _handleReadyPosition(List<Object?> event) {
    final poseLandmarksBuffered = event.first as List<List<PoseLandmark>>;
    final stepExercise = event.last as StepExercise;
    if (stepExercise != StepExercise.notInPlace) return;
    for (final poseLandmarks in poseLandmarksBuffered) {
      for (final poseLandmark in poseLandmarks) {
        if (poseLandmark.inFrameLikelihood < 0.8) return;
      }
    }
    _stepExerciseController.add(StepExercise.ready);
  }

  Stream<List<PoseLandmark>> _getExerciseJointsStream(Stream<Pose> stream) {
    return stream
        .where((pose) => pose.landmarks.isNotEmpty)
        .map((pose) => pose.landmarks.where((landmark) => Exercise.authorizedType.contains(landmark.type)).toList());
  }

  Stream<MeanFilteredData> _getMeanFilterStream(Stream<List<PoseLandmark>> stream) {
    return stream
        // Get last [buffer] poses
        .bufferCount(meanFilterBuffer, 1)
        // Swap matrix [buffer] * [authorizedType.length]
        .map(_swapMatrixDimensions)
        // For every landmarks, get meanFilter of size [buffer]
        .map((filteredLandmarks) => filteredLandmarks.map(_meanFilter).toList());
  }

  List<double> _meanFilter(List<PoseLandmark> landmarks) {
    return landmarks
        .map((landmark) => landmark.position)
        .map((position) => [
              position.x / meanFilterBuffer,
              position.y / meanFilterBuffer,
              position.z / meanFilterBuffer,
            ])
        .reduce((value, element) => [
              value[0] + element[0],
              value[1] + element[1],
              value[2] + element[2],
            ]);
  }

  List<List<T>> _swapMatrixDimensions<T>(List<List<T>> matrix) {
    final height = matrix.length;
    final width = matrix[0].length;
    List<List<T>> newMatrix = [];
    for (int col = 0; col < width; col++) {
      List<T> newRow = [];
      for (int row = 0; row < height; row++) {
        newRow.add(matrix[row][col]);
      }
      newMatrix.add(newRow);
    }
    return newMatrix;
  }

  Future<void> _startCameraStream() async {
    await BodyDetection.startCameraStream(onFrameAvailable: _handleCameraImage, onPoseAvailable: _handlePose);
    await BodyDetection.enablePoseDetection();
  }

  Future<void> _stopCameraStream() async {
    await BodyDetection.disablePoseDetection();
    await BodyDetection.stopCameraStream();
  }

  void _handleCameraImage(ImageResult result) {
    if (!mounted) return;

    // To avoid a memory leak issue.
    // https://github.com/flutter/flutter/issues/60160
    PaintingBinding.instance?.imageCache?.clear();
    PaintingBinding.instance?.imageCache?.clearLiveImages();

    final image = Image.memory(
      result.bytes,
      gaplessPlayback: true,
      fit: BoxFit.contain,
    );

    setState(() {
      _cameraImage = image;
      _imageSize = result.size;
    });
  }

  void _handlePose(Pose? pose) {
    if (!mounted) return;
    if (pose != null) {
      _poseController.add(pose);
    }
    setState(() {
      _detectedPose = pose;
    });
  }

  @override
  void dispose() {
    _stopCameraStream();
    for (final ss in _streamSubscriptions) {
      ss.cancel();
    }
    for (final sc in _streamControllers) {
      sc.close();
    }
    super.dispose();
  }

  @override
  Widget build(BuildContext context) {
    return FutureBuilder<void>(
      future: _startCamera,
      builder: (context, snapshot) {
        if (snapshot.connectionState == ConnectionState.waiting) {
          return const Center(child: CircularProgressIndicator());
        }
        return Column(
          children: [
            Center(
              child: CustomPaint(
                child: _cameraImage,
                // foregroundPainter: PosePainter(
                //   pose: _detectedPose,
                //   imageSize: _imageSize,
                // ),
              ),
            ),
            ExerciseIndicator(
              stepExerciseController: _stepExerciseController,
              repCounter: _repCounter,
              exercise: widget.exercise,
            ),
          ],
        );
      },
    );
  }
}