unkraut/main.py at main · keimschleuder/unkraut · 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
102
103
104
105
106
107
108
109
110
111
112
from picamera2 import Picamera2
import io
import os
from PIL import Image
import time
import argparse
import rpi_drivers.arduino
import rpi_drivers.plantnet
import rpi_drivers.servo

# Console Arguments init
parser = argparse.ArgumentParser(prog = "main.py",
                                 usage = "python main.py [-h]",
                                 description = "Intended to run on a Raspberry Pi 5 with AI-HAT and AI-Camera. It detects harmful plants and destroys them.",
                                 epilog = "Program written for the german Jugend Forscht-Competition",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

# Driver init
arduino_drv = rpi_drivers.arduino.Arduino()

plantnet_drv = rpi_drivers.plantnet.PlantNet()

roboterarm = rpi_drivers.servo.Servo()

def find_weeds(bestMatches):
    # Evaluate whether a crop is a weed based on a list in the plantnet software driver
    weeds_found = []
    for i, match in enumerate(bestMatches):
        if match is not None:
            for weed in rpi_drivers.plantnet.weeds:
                if match.lower().find(weed) != -1:
                    weeds_found.append([weed, i])
    return weeds_found

def treat_weeds(weeds_found):
    for weed, quadrant in weeds_found:
        print(f"Weed '{weed}' found in quadrant {quadrant + 1}. Treating...")
        match quadrant:
            case 0:
                roboterarm.goto(0, 45)
                roboterarm.goto(1, 180)
                roboterarm.goto(2, 100)
            case 1:
                roboterarm.goto(0, 0)
                roboterarm.goto(1, 180)
                roboterarm.goto(2, 100)
            case 2:
                roboterarm.goto(0, 45)
                roboterarm.goto(1, 180)
                roboterarm.goto(2, 10)
            case 3:
                roboterarm.goto(0, 0)
                roboterarm.goto(1, 180)
                roboterarm.goto(2, 10)

        arduino_drv.pump(True)
        time.sleep(5)
        arduino_drv.pump(False)

def main():
    # Start Camera
    picam2 = Picamera2()
    picam2.start()
    time.sleep(1)

    for i in range(5):
        # Read image
        data = io.BytesIO()
        picam2.capture_file(data, format='jpeg')

        data.seek(0)
        img = Image.open(data).convert("RGB")
        w, h = img.size
        mx, my = w // 2, h // 2

        # Split into quadrants
        q1 = img.crop((0, 0, mx, my))       # top-left
        q2 = img.crop((mx, 0, w, my))       # top-right
        q3 = img.crop((0, my, mx, h))       # bottom-left
        q4 = img.crop((mx, my, w, h))       # bottom-right

        quadrants = []
        for q in (q1, q2, q3, q4):
            buf = io.BytesIO()
            q.save(buf, format='JPEG')
            buf.seek(0)
            quadrants.append(buf)

        '''
        # Debugging: Save quadrants to files
        output_dir = os.path.join(os.path.dirname(__file__), "captures")
        ts = int(time.time() * 1000)
        for i, q in enumerate((q1, q2, q3, q4), start=1):
            fn = os.path.join(output_dir, f"{ts}_q{i}.jpg")
            q.save(fn, format="JPEG")
        print("Saved quadrants.")
        '''

        # Send the quadrants to the API and process the results
        bestMatches = plantnet_drv.send_multiple_requests(quadrants)
        weeds_found = find_weeds(bestMatches)
        if len(weeds_found) > 0:
            treat_weeds(weeds_found)
            print("All weeds treated")
        else:
            print("No weeds found in this capture.")

        arduino_drv.motor(True)
        time.sleep(1)
        arduino_drv.motor(False)

main()