Project-TALON/animal_finder.py at main · I-am-Chittesh/Project-TALON · 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
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
89
90
91
92
93
94
95
96
97
98
99
100
101
import cv2
import cvlib as cv
import numpy as np
import urllib.request
import time

def main():
    print("--- AI Animal Finder ---")


    esp32_ip = input("Enter the ESP32-CAM's IP address (e.g., 192.168.4.1): ")
    stream_url = f'http://{esp32_ip}:81/stream'


    target_animal = input("Which animal do you want to find? (e.g., 'dog', 'cat', 'bird'): ").lower()

    print(f"\n[INFO] Starting stream from: {stream_url}")
    print(f"[INFO] Looking for: {target_animal}")
    print("[INFO] Press 'q' to quit.")

    stream = None
    try:

        stream = urllib.request.urlopen(stream_url)
    except Exception as e:
        print(f"\n[ERROR] Could not open stream: {e}")
        print("Please check the IP address and make sure the robot is on.")
        return

    bytes_buffer = b''


    while True:
        try:

            bytes_buffer += stream.read(1024)


            a = bytes_buffer.find(b'\xff\xd8')
            b = bytes_buffer.find(b'\xff\xd9')


            if a != -1 and b != -1:

                jpg = bytes_buffer[a:b+2]

                bytes_buffer = bytes_buffer[b+2:]


                frame = cv2.imdecode(np.frombuffer(jpg, dtype=np.uint8), cv2.IMREAD_COLOR)

                if frame is None:
                    print("[WARN] Received empty frame. Skipping.")
                    continue


                bbox, labels, conf = cv.detect_common_objects(frame, model='yolov4-tiny')

                found_target = False


                for label, c in zip(labels, conf):
                    if label == target_animal:
                        print(f"!!! TARGET FOUND: {label} (Confidence: {c*100:.2f}%) !!!")
                        found_target = True


                output_frame = cv.object_detection.draw_bbox(frame, bbox, labels, conf, write_conf=True)

                if found_target:
                    cv2.putText(output_frame, "!!! TARGET ACQUIRED !!!", (10, 30),
                                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 3)


                cv2.imshow("AI Animal Finder", output_frame)


                if cv2.waitKey(1) & 0xFF == ord('q'):
                    print("[INFO] Quitting...")
                    break

        except Exception as e:
            print(f"[ERROR] Stream error: {e}")
            time.sleep(2)


            try:
                stream.close()
                stream = urllib.request.urlopen(stream_url)
                bytes_buffer = b''
            except Exception as e:
                print(f"[ERROR] Reconnect failed: {e}")
                time.sleep(5)


    cv2.destroyAllWindows()
    if stream:
        stream.close()

if __name__ == '__main__':
    main()