DNP/predictBoundingBox.py at main · hoangvanphi2004/DNP · GitHub

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
import consumer
import producer
import boundingBoxDetection
import cv2 as cv
import os

if __name__ == "__main__":
    latestFrameConsumer = consumer.LatestFrameConsumer("latest_frame")
    frameProducer = producer.FrameProducer()
    boundingBoxPrediction = boundingBoxDetection.YOLOv8()
    boundingBoxProducer = producer.BoundingBoxProducer()
    keypointsConsumer = consumer.KeypointsConsumer("receive_back_keypoints")
    setupProducer = producer.SetupProducer()
    setupProducer.send_setup_value(0)
    setupProducer.producer.flush()
    setupConsumer = consumer.SetupConsumer("end")

    boundingBox = None
    firstTime = True
    isEnd = False
    try:
        while True:
            retKeypoints, keypointsData, _ = keypointsConsumer.receive_keypoints()
            retSetup, value = setupConsumer.receive_setup_value()

            if retSetup and value == 3:
                isEnd = True

            if boundingBox == None:
                retFrame, frame, offset = latestFrameConsumer.receive_latest_frame()
                if retFrame:
                    boundingBox = boundingBoxPrediction.predict(frame)
                    boundingBox = {
                        "offset": offset,
                        "data": boundingBox,
                    }
                elif isEnd:
                    setupProducer.send_setup_value(5)
                    break

            if (retKeypoints or firstTime) and boundingBox != None:
                frameProducer.send_frame(frame, offset)
                boundingBoxProducer.send_bounding_box(boundingBox)
                boundingBox = None
                firstTime = False
    except KeyboardInterrupt:
        pass
    finally:
        latestFrameConsumer.consumer.close()
        keypointsConsumer.consumer.close()
        boundingBoxProducer.producer.flush()
        frameProducer.producer.flush()
        setupConsumer.consumer.close()