Skip to content

Integrate CKN with MegaDetector #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

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.

1 participant