Ading data collection script for spell capturing

collect_spell_data.py

@@ -0,0 +1,77 @@
+from KanoWandAsync import Shop, Wand
+from KanoWandAsync.constants import *
+import asyncio
+import pandas as pd
+from uuid import uuid4
+from pathlib import Path
+import os
+global spell
+def save_dataframe(data):
+    df = pd.DataFrame(data)
+    df.columns = ["x", "y", "pitch", "roll"]
+    path = Path(f"./data/{spell}/{uuid4()}.csv")
+    if not path.parent.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+    df.to_csv(str(path), index=False)
+
+
+class DataCollectWand(Wand):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.colors = ["#a333c8", "2185d0", "0x21ba45", "#fbbd08", "#f2711c", "#db2828"]
+        self.position_id = None
+        self.data = []
+
+    async def post_connect(self):
+        print("Connected to {}".format(self.name))
+        await self.subscribe_button()
+
+    async def on_button(self, pressed):
+        if pressed:
+            self.data = []
+            print("Button pressed")
+            await self.vibrate(PATTERN.BURST)
+            await self.reset_position()
+            await self.subscribe_position()
+        else:
+            save_dataframe(self.data)
+            await self.unsubscribe_position()
+
+    async def on_position(self, x, y, pitch, roll):
+        # print(x, y, pitch, roll)
+        self.data.append([x, y, pitch, roll])
+
+def get_prompt_text(spells):
+    strings = [f"{spells[spells == spell].index[0]} - {spell} " for spell in spells]
+    return "Enter number to select spell- \n" + "\n".join(strings)
+
+async def main():
+    spells = pd.read_csv("spells.csv")["spells"]
+    shop = Shop(asyncio.get_running_loop(), wand_class=DataCollectWand, debug=True)
+    wands = []
+    x = input(get_prompt_text(spells))
+
+    while x != "quit":
+        try:
+            global spell
+            spell = spells[int(x)]
+            print(f"Currently chosen spell - {spell}")
+        except Exception as e:
+            print(e)
+
+        try:
+            while len(wands) == 0:
+                print("Scanning...")
+                wands = await shop.scan(connect=True)
+                for wand in wands:
+                    await wand.vibrate(PATTERN.BURST)
+                    await asyncio.sleep(60)
+        except Exception as e:
+            print(e)
+            for wand in wands:
+                wand.disconnect()
+        x = input(get_prompt_text(spells))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())