Skip to content

Conversation

@nee1k
Copy link
Collaborator

@nee1k nee1k commented Feb 14, 2025

This pull request introduces a new model class MegaDetectorModel in the plugins/ckn_inference_daemon/models/mega_detector.py file. The class is designed to load the MegaDetector V5 checkpoint from Hugging Face, preprocess input images, run inference, and return detections.

Key changes include:

New Model Class Implementation:

  • Class Definition and Initialization:

    • Introduced the MegaDetectorModel class, inheriting from BaseModel. This class initializes with parameters for device type, confidence threshold, IoU threshold, and an optional labels path. It also handles loading the MegaDetector V5 checkpoint from Hugging Face and setting up the model for inference.
  • Image Preprocessing:

    • Added the pre_process method to handle image loading, conversion to RGB, letterboxing, and tensor conversion. This method prepares the image for inference by the model.
  • Inference and Prediction:

    • Implemented the predict method to run the model on pre-processed input data, apply non-max suppression, and rescale detections to the original image size. This method returns a list of detections with labels and confidence scores.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants