| 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 | |