Upload 11 files

yoloxdetect2/configs/yolov3.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 1.0
+        self.width = 1.0
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self, sublinear=False):
+        def init_yolo(M):
+            for m in M.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eps = 1e-3
+                    m.momentum = 0.03
+        if "model" not in self.__dict__:
+            from yolox.models import YOLOX, YOLOFPN, YOLOXHead
+            backbone = YOLOFPN()
+            head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act="lrelu")
+            self.model = YOLOX(backbone, head)
+        self.model.apply(init_yolo)
+        self.model.head.initialize_biases(1e-2)
+
+        return self.model

yoloxdetect2/configs/yolox_l.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 1.0
+        self.width = 1.0
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

yoloxdetect2/configs/yolox_m.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.67
+        self.width = 0.75
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

yoloxdetect2/configs/yolox_nano.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.25
+        self.input_size = (416, 416)
+        self.random_size = (10, 20)
+        self.mosaic_scale = (0.5, 1.5)
+        self.test_size = (416, 416)
+        self.mosaic_prob = 0.5
+        self.enable_mixup = False
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self, sublinear=False):
+
+        def init_yolo(M):
+            for m in M.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eps = 1e-3
+                    m.momentum = 0.03
+        if "model" not in self.__dict__:
+            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+            in_channels = [256, 512, 1024]
+            # NANO model use depthwise = True, which is main difference.
+            backbone = YOLOPAFPN(
+                self.depth, self.width, in_channels=in_channels,
+                act=self.act, depthwise=True,
+            )
+            head = YOLOXHead(
+                self.num_classes, self.width, in_channels=in_channels,
+                act=self.act, depthwise=True
+            )
+            self.model = YOLOX(backbone, head)
+
+        self.model.apply(init_yolo)
+        self.model.head.initialize_biases(1e-2)
+        return self.model

yoloxdetect2/configs/yolox_s.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.50
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

yoloxdetect2/configs/yolox_tiny.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.375
+        self.input_size = (416, 416)
+        self.mosaic_scale = (0.5, 1.5)
+        self.random_size = (10, 20)
+        self.test_size = (416, 416)
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+        self.enable_mixup = False

yoloxdetect2/configs/yolox_x.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 1.33
+        self.width = 1.25
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

yoloxdetect2/helpers.py

@@ -0,0 +1,105 @@
+from yoloxdetect.utils.downloads import attempt_download_from_hub, attempt_download
+from yolox.data.datasets import COCO_CLASSES
+from yolox.data.data_augment import preproc
+from yolox.utils import postprocess, vis
+import importlib
+import torch
+import cv2
+import os
+
+
+class YoloxDetector2:
+    def __init__(
+        self,
+        model_path: str,
+        config_path: str,
+        device: str = "cpu",
+        hf_model: bool = False,
+    ):
+
+        self.device = device
+        self.config_path = config_path
+        self.classes = COCO_CLASSES
+        self.conf = 0.3
+        self.iou = 0.45
+        self.show = False  
+        self.save = True
+        self.torchyolo = False
+        
+        if self.save:
+            self.save_path = 'output/result.jpg'
+            
+        if hf_model:
+            self.model_path = attempt_download_from_hub(model_path)
+        
+        else:
+            self.model_path = attempt_download(model_path)        
+    
+        self.load_model()
+
+
+    def load_model(self):
+        current_exp = importlib.import_module(self.config_path)
+        exp = current_exp.Exp()
+        
+        model = exp.get_model()
+        model.to(self.device)
+        model.eval()
+        ckpt = torch.load(self.model_path, map_location=self.device)
+        model.load_state_dict(ckpt["model"])
+        self.model = model
+
+
+    def predict(self, image_path, image_size):    
+        image = cv2.imread(image_path)    
+        if image_size is not None:
+            ratio = min(image_size / image.shape[0], image_size / image.shape[1])
+            img, _ = preproc(image, input_size=(image_size, image_size))
+            img = torch.from_numpy(img).to(self.device).unsqueeze(0).float()
+        else:
+            manuel_size = 640
+            ratio = min(manuel_size / image.shape[0], manuel_size / image.shape[1])
+            img, _ = preproc(image, input_size=(manuel_size, manuel_size))
+            img = torch.from_numpy(img).to(self.device).unsqueeze(0).float()
+
+        prediction_result = self.model(img)
+        original_predictions = postprocess(
+            prediction=prediction_result,
+            num_classes= len(COCO_CLASSES),
+            conf_thre=self.conf,
+            nms_thre=self.iou)[0]
+        
+        if original_predictions is None :
+           return None
+        output = original_predictions.cpu()
+        bboxes = output[:, 0:4]
+        bboxes /= ratio
+        cls = output[:, 6]
+        scores = output[:, 4] * output[:, 5]
+        if self.torchyolo is False:
+            vis_res = vis(
+                image,
+                bboxes,
+                scores,
+                cls,
+                self.conf,
+                COCO_CLASSES,
+            )
+            if self.show:
+                cv2.imshow("result", vis_res)
+                cv2.waitKey(0)
+                cv2.destroyAllWindows()
+            elif self.save:
+                save_dir = self.save_path[:self.save_path.rfind('/')]
+                if not os.path.exists(save_dir):
+                    os.makedirs(save_dir)
+                cv2.imwrite(self.save_path, vis_res)
+                return self.save_path
+            
+            else:
+                return vis_res
+        else:
+            object_predictions_list = [bboxes, scores, cls, COCO_CLASSES]
+            return object_predictions_list
+
+

yoloxdetect2/utils/downloads.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+import torch
+import urllib
+import requests
+import subprocess
+
+def attempt_download_from_hub(repo_id, hf_token=None):
+    # https://github.com/fcakyon/yolov5-pip/blob/main/yolov5/utils/downloads.py
+    from huggingface_hub import hf_hub_download, list_repo_files
+    from huggingface_hub.utils._errors import RepositoryNotFoundError
+    from huggingface_hub.utils._validators import HFValidationError
+    try:
+        repo_files = list_repo_files(repo_id=repo_id, repo_type='model', token=hf_token)
+        model_file = [f for f in repo_files if f.endswith('.pth')][0]
+        file = hf_hub_download(
+            repo_id=repo_id,
+            filename=model_file,
+            repo_type='model',
+            token=hf_token,
+        )
+        return file
+    except (RepositoryNotFoundError, HFValidationError):
+        return None
+    
+    
+def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
+    import os
+    # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
+
+    file = Path(file)
+    assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
+    try:  # url1
+        torch.hub.download_url_to_file(url, str(file), progress=True)  # pytorch download
+        assert file.exists() and file.stat().st_size > min_bytes, assert_msg  # check
+    except Exception as e:  # url2
+        file.unlink(missing_ok=True)  # remove partial downloads
+        os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -")  # curl download, retry and resume on fail
+    finally:
+        if not file.exists() or file.stat().st_size < min_bytes:  # check
+            file.unlink(missing_ok=True)  # remove partial downloads
+            raise Exception(error_msg or assert_msg)  # raise informative error
+
+def attempt_download(file, repo='Megvii-BaseDetection/YOLOX', release='0.1.0'):
+    def github_assets(repository, version='latest'):
+        response = requests.get(f'https://api.github.com/repos/{repository}/releases/tags/{version}').json()  # github api
+        return response['tag_name'], [x['name'] for x in response['assets']]  # tag, assets
+
+    file = Path(str(file).strip().replace("'", ''))
+    if not file.exists():
+        # URL specified
+        name = Path(urllib.parse.unquote(str(file))).name  # decode '%2F' to '/' etc.
+        if str(file).startswith(('http:/', 'https:/')):  # download
+            url = str(file).replace(':/', '://')  # Pathlib turns :// -> :/
+            file = name.split('?')[0]  # parse authentication https://url.com/file.txt?auth...
+            if Path(file).is_file():
+                return file
+            else:
+                safe_download(file=file, url=url, min_bytes=1E5)
+            return file
+
+        # GitHub assets
+        assets = [
+            'yolov6n.pt', 'yolov6s.pt', 'yolov6m.pt', 'yolov6l.pt',
+            'yolov6n6.pt', 'yolov6s6.pt', 'yolov6m6.pt', 'yolov6l6.pt']
+        try:
+            tag, assets = github_assets(repo, release)
+        except Exception:
+            try:
+                tag, assets = github_assets(repo)  # latest release
+            except Exception:
+                try:
+                    tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
+                except Exception:
+                    tag = release
+
+        file.parent.mkdir(parents=True, exist_ok=True)  # make parent dir (if required)
+        if name in assets:
+            safe_download(
+                file,
+                url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
+                url2=f'https://storage.googleapis.com/{repo}/{tag}/{name}',  # backup url (optional)
+                min_bytes=1E5,
+                error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/{tag}')
+    
+    return str(file)