fixing plant_locator

zauberzeug · Nov 25, 2024 · 38aaaf7 · 38aaaf7
1 parent 340073f
commit 38aaaf7
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/field_friend/automations/plant_locator.py b/field_friend/automations/plant_locator.py
@@ -10,7 +10,7 @@
 from .plant import Plant
 
 WEED_CATEGORY_NAME = ['coin', 'weed', 'weedy_area', ]
-CROP_CATEGORY_NAME = {}
+CROP_CATEGORY_NAME: dict[str, str] = {}
 MINIMUM_CROP_CONFIDENCE = 0.3
 MINIMUM_WEED_CONFIDENCE = 0.3
 
@@ -38,7 +38,7 @@ def __init__(self, system: 'System') -> None:
         self.autoupload: Autoupload = Autoupload.DISABLED
         self.upload_images: bool = False
         self.weed_category_names: list[str] = WEED_CATEGORY_NAME
-        self.crop_category_names: list[str] = CROP_CATEGORY_NAME
+        self.crop_category_names: dict[str, str] = CROP_CATEGORY_NAME
         self.minimum_crop_confidence: float = MINIMUM_CROP_CONFIDENCE
         self.minimum_weed_confidence: float = MINIMUM_WEED_CONFIDENCE
         rosys.on_repeat(self._detect_plants, 0.01)  # as fast as possible, function will sleep if necessary
@@ -127,10 +127,10 @@ async def _detect_plants(self) -> None:
             if d.category_name in self.weed_category_names and d.confidence >= self.minimum_weed_confidence:
                 # self.log.info('weed found')
                 await self.plant_provider.add_weed(plant)
-            elif d.category_name in self.crop_category_names and d.confidence >= self.minimum_crop_confidence:
+            elif d.category_name in self.crop_category_names.keys() and d.confidence >= self.minimum_crop_confidence:
                 # self.log.info(f'{d.category_name} found')
                 self.plant_provider.add_crop(plant)
-            elif d.category_name not in self.crop_category_names and d.category_name not in self.weed_category_names:
+            elif d.category_name not in self.crop_category_names.keys() and d.category_name not in self.weed_category_names:
                 self.log.info(f'{d.category_name} not in categories')
             # else:
             #     self.log.info(f'confidence of {d.category_name} to low: {d.confidence}')
@@ -231,10 +231,12 @@ async def get_crop_names(self) -> dict[str, str]:
             for weed in weeds:
                 crop_names.remove(weed)
             CROP_CATEGORY_NAME.update({name: name.replace('_', ' ').title() for name in crop_names})
+            self.crop_category_names = CROP_CATEGORY_NAME
             return CROP_CATEGORY_NAME
         else:
             simulated_crop_names: list[str] = ['coin_with_hole', 'borrietsch', 'estragon', 'feldsalat', 'garlic', 'jasione', 'kohlrabi', 'liebstoeckel', 'maize', 'minze', 'onion',
                                                'oregano_majoran', 'pastinake', 'petersilie', 'pimpinelle', 'red_beet', 'salatkopf', 'schnittlauch', 'sugar_beet', 'thymian_bohnenkraut', 'zitronenmelisse', ]
 
             CROP_CATEGORY_NAME.update({name: name.replace('_', ' ').title() for name in simulated_crop_names})
+            self.crop_category_names = CROP_CATEGORY_NAME
             return CROP_CATEGORY_NAME