Added pip support.

README.md

@@ -8,19 +8,25 @@
 
 The TorchYolo library aims to support all YOLO models(like YOLOv5, YOLOv6, YOLOv7, YOLOX, etc.) and provide a unified interface for training and inference. The library is based on PyTorch and is designed to be easy to use and extend.
 
-### Installation
-```bash
-git clone https://github.com/kadirnar/torchyolo
-cd torchyolo
-pip install -r requirements.txt
-```
-or 
+### Installation 
 ```bash
 pip install torchyolo
 ```
 ### Usage
 ```python
-python torchyolo/predict.py --config configs/default_config.yaml
+from torchyolo import YoloPredictor
+predictor = YoloPredictor(
+  model_type="yolov5", 
+  model_path="yolov5s.pt", 
+  device='cpu', 
+  image_size=640
+)
+predictor.conf_thres = 0.25
+predictor.iou_thres = 0.45
+predictor.save = True
+predictor.show = False
+image = "data/highway.jpg"
+result = predictor.predict(image)
 ```
 Note: You only need to make changes in the default_config.yaml file.
 
@@ -40,6 +46,9 @@ Before opening a PR:
 - [ ] Add Train, Export and Eval scripts
 - [ ] Add Benchmark Results
 
+### Acknowledgement
+A part of the code is borrowed from [SAHI](https://github.com/obss/sahi). Many thanks for their wonderful works.
+
 ### Citation
 ```bibtex
 @article{li2022yolov6,
@@ -68,13 +77,13 @@ Before opening a PR:
 ```bibtex
 @software{glenn_jocher_2020_4154370,
   author       = {Glenn Jocher and,Alex Stoken and,Jirka Borovec and,NanoCode012 and,ChristopherSTAN and,Liu Changyu and,Laughing and,tkianai and,Adam Hogan and,lorenzomammana and,yxNONG and,AlexWang1900 and,Laurentiu Diaconu and,Marc and,wanghaoyang0106 and,ml5ah and,Doug and,Francisco Ingham and,Frederik and,Guilhen and,Hatovix and,Jake Poznanski and,Jiacong Fang and,Lijun Yu 于力军 and,changyu98 and,Mingyu Wang and,Naman Gupta and,Osama Akhtar and,PetrDvoracek and,Prashant Rai},
-  title        = {{ultralytics/yolov5: v7.2 - Bug Fixes and 
+  title={{ultralytics/yolov5: v7.2 - Bug Fixes and 
                    Performance Improvements}},
-  month        = oct,
-  year         = 2020,
-  publisher    = {Zenodo},
-  version      = {v3.1},
-  doi          = {10.5281/zenodo.4154370},
-  url          = {https://doi.org/10.5281/zenodo.4154370}
+  month= oct,
+  year= 2020,
+  publisher= {Zenodo},
+  version= {v3.1},
+  doi= {10.5281/zenodo.4154370},
+  url= {https://doi.org/10.5281/zenodo.4154370}
 }
 ```

torchyolo/__init__.py

@@ -1 +1,3 @@
-__version__ = "0.1.2"
+from torchyolo.predict import YoloPredictor
+
+__version__ = "0.1.3"

torchyolo/predict.py

@@ -1,29 +1,40 @@
-from automodel import AutoDetectionModel
+from torchyolo.automodel import AutoDetectionModel
 
-from torchyolo.utils.config_utils import get_config
 
+class YoloPredictor:
+    def __init__(self, model_type="yolov5", model_path="yolov5s.pt", device="cpu", image_size=640):
+        self.model_type = model_type
+        self.model_path = model_path
+        self.config_path = "torchyolo.configs.yolox.yolox_m"  # yolox_nano.py
+        self.device = device
+        self.conf_thres = 0.25
+        self.iou_thres = 0.45
+        self.image_size = image_size
+        self.save = True
+        self.show = False
 
-def model_prediction(config_file):
-    config = get_config(config_file)
-    model = AutoDetectionModel.from_pretrained(
-        model_type=config.DETECTOR_CONFIG.MODEL_TYPE,
-        model_path=config.DETECTOR_CONFIG.MODEL_PATH,
-        config_path=config.DETECTOR_CONFIG.CONFIG_PATH,
-        device=config.DETECTOR_CONFIG.DEVICE,
-        confidence_threshold=config.DETECTOR_CONFIG.CONFIDENCE_THRESHOLD,
-        iou_threshold=config.DETECTOR_CONFIG.IOU_THRESHOLD,
-        image_size=config.DETECTOR_CONFIG.IMAGE_SIZE,
-    )
-    model.show = config.DETECTOR_CONFIG.SHOW
-    model.save = config.DETECTOR_CONFIG.SAVE
+        # Load Model
+        self.load_model()
 
-    model.predict(image=config.FILE_CONFIG.IMAGE_PATH)
+    def load_model(self):
+        model = AutoDetectionModel.from_pretrained(
+            model_type=self.model_type,
+            model_path=self.model_path,
+            config_path=self.config_path,
+            device=self.device,
+            confidence_threshold=self.conf_thres,
+            iou_threshold=self.iou_thres,
+            image_size=self.image_size,
+        )
+        model.save = self.save
+        model.show = self.show
+        self.model = model
 
+    def predict(self, image):
+        return self.model.predict(image)
 
-if __name__ == "__main__":
-    import argparse
 
-    args = argparse.ArgumentParser()
-    args.add_argument("--config_file", type=str, default="torchyolo/configs/default_config.yaml")
-    args = args.parse_args()
-    model_prediction(args.config_file)
+if __name__ == "__main__":
+    predictor = YoloPredictor(model_type="yolov5", model_path="yolov5s.pt", device="cpu", image_size=640)
+    image = "data/highway.jpg"
+    result = predictor.predict(image)