File size: 3,312 Bytes
c36e146 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
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
|