Upload mega_detector.py

mega_detector.py

@@ -10,22 +10,22 @@ from abc import ABC, abstractmethod
 class BaseModel(ABC):
     @abstractmethod
     def pre_process(self, filename: str):
-        """
-        Pre-process the input file and return it as a tensor.
-        """
+        """Pre-process the input file and return it as a tensor."""
         pass
 
     @abstractmethod
     def predict(self, input_data):
-        """
-        Run inference on the pre-processed input and return predictions.
-        """
+        """Run inference on the pre-processed input and return predictions."""
         pass
 
 class MegaDetectorModel(BaseModel):
     """
-    MegaDetectorModel loads the MegaDetector V5 checkpoint from Hugging Face,
+    MegaDetectorModel loads the MegaDetector checkpoint from a Hugging Face repository,
     preprocesses input images, runs inference, and returns detections (label/confidence).
+
+    The repository ID is the only input required. The model filename, class name, and weight file
+    are all expected to match the repository's base name. For example, if the repository ID is
+    "nkarthikeyan/MegaDetectorV5", then the model weight file should be "MegaDetectorV5.pt".
     """
 
     def __init__(self, device='cpu', conf_thres=0.25, iou_thres=0.45, labels_path=None):
@@ -41,8 +41,11 @@ class MegaDetectorModel(BaseModel):
     @classmethod
     def from_pretrained(cls, repo_id: str, device: str = 'cpu', **kwargs):
         """
-        Loads the model checkpoint from the given Hugging Face repository and returns
-        an instance of MegaDetectorModel ready for inference.
+        Loads the model checkpoint from the given Hugging Face repository and returns an instance
+        of MegaDetectorModel ready for inference.
+
+        The repository's base name is used to derive the model weight filename. For example, if
+        repo_id is "nkarthikeyan/MegaDetectorV5", then the weight file is expected to be "MegaDetectorV5.pt".
 
         Args:
             repo_id (str): The Hugging Face repository ID (e.g. "nkarthikeyan/MegaDetectorV5").
@@ -52,8 +55,10 @@ class MegaDetectorModel(BaseModel):
             MegaDetectorModel: An instance with the model loaded.
         """
         instance = cls(device=device, **kwargs)
-        # Download the model checkpoint (assumes the file is named 'model.pt')
-        model_path = hf_hub_download(repo_id=repo_id, filename="model.pt")
+        # Use the repository base name as the weight filename.
+        model_name = repo_id.split("/")[-1]
+        weight_filename = f"{model_name}.pt"
+        model_path = hf_hub_download(repo_id=repo_id, filename=weight_filename)
         checkpoint = torch.load(model_path, map_location=instance.device)
         instance.model = checkpoint['model'].float().fuse().eval()
         if instance.device.type != 'cpu':