Skip to content

Commit

Permalink
removes PIL, adds two new methods
Browse files Browse the repository at this point in the history
  • Loading branch information
bhaney committed May 21, 2024
1 parent 66b3338 commit d9bd30f
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 9 deletions.
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
viam-sdk >= 0.5.1
viam-sdk >= 0.21.0
pillow
52 changes: 44 additions & 8 deletions src/verificationclassifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,16 @@

from enum import Enum

from PIL import Image

from viam.media.video import RawImage
from viam.media.video import ViamImage

from viam.module.types import Reconfigurable
from viam.proto.app.robot import ComponentConfig
from viam.proto.service.vision import Classification
from viam.proto.common import ResourceName
from viam.proto.service.vision import Classification, Detection
from viam.proto.common import ResourceName, PointCloudObject
from viam.resource.base import ResourceBase
from viam.resource.types import Model, ModelFamily

from viam.services.vision import Vision
from viam.services.vision import Vision, CaptureAllResult

from viam.components.camera import Camera

Expand Down Expand Up @@ -162,6 +160,44 @@ async def do_command(self):
async def get_object_point_clouds(self):
return

async def get_properties(
self,
*,
extra: Optional[Mapping[str, Any]] = None,
timeout: Optional[float] = None) -> Vision.Properties:
return Vision.Properties(
classifications_supported=True,
detections_supported=False,
object_point_clouds_supported=False,
)

async def capture_all_from_camera(
self,
camera_name: str,
return_image: bool = False,
return_classifications: bool = False,
return_detections: bool = False,
return_object_point_clouds: bool = False,
*,
extra: Optional[Mapping[str, Any]] = None,
timeout: Optional[float] = None,
) -> CaptureAllResult:
result = CaptureAllResult()
if camera_name != self.camera_name:
raise Exception(
f"camera {camera_name} was not declared in the camera_name dependency")
cam_image = await self.camera.get_image(mime_type="image/jpeg")
if return_image:
result.image = cam_image
if return_classifications:
cls = await self.get_classifications(cam_image, 1)
result.classifications = cls
if return_detections:
result.detections = []
if return_object_point_clouds:
result.objects = []
return result

async def get_classifications_from_camera(self,
camera_name: str,
count: int,
Expand All @@ -177,7 +213,7 @@ async def get_classifications_from_camera(self,
return await self.get_classifications(cam_image, 1)

async def get_classifications(self,
image: Union[Image.Image, RawImage],
image: ViamImage,
count: int,
*,
extra: Optional[Dict[str, Any]] = None,
Expand All @@ -199,7 +235,7 @@ async def get_classifications(self,
classifications = [{"class_name": class_name, "confidence": 1.0}]
return classifications

async def process_image(self, image: Union[Image.Image, RawImage]):
async def process_image(self, image: ViamImage):
if self.alarm_state is AlarmState.TRIGGER_1:
if self.trigger_1_detector is None:
self.alarm_state = AlarmState.TRIGGER_2 # go straight to trigger 2
Expand Down

0 comments on commit d9bd30f

Please sign in to comment.