Spaces:
Runtime error
Runtime error
File size: 2,447 Bytes
37f5c2f |
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 |
import os
import time
from loguru import logger
import cv2
import torch
from yolox.data.data_augment import ValTransform
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import postprocess, vis
class Predictor(object):
def __init__(
self,
model,
cls_names=COCO_CLASSES,
device="cpu",
fp16=False,
legacy=False,
):
self.model = model
self.cls_names = cls_names
self.num_classes = len(COCO_CLASSES)
self.confthre = 0.01
self.nmsthre = 0.01
self.test_size = (640, 640)
self.device = device
self.fp16 = fp16
self.preproc = ValTransform(legacy=legacy)
def inference(self, img, confthre=None, nmsthre=None, test_size=None):
if confthre is not None:
self.confthre = confthre
if nmsthre is not None:
self.nmsthre = nmsthre
if test_size is not None:
self.test_size = test_size
img_info = {"id": 0}
if isinstance(img, str):
img_info["file_name"] = os.path.basename(img)
img = cv2.imread(img)
else:
img_info["file_name"] = None
cv2.imwrite("test.png", img)
height, width = img.shape[:2]
img_info["height"] = height
img_info["width"] = width
img_info["raw_img"] = img
ratio = min(self.test_size[0] / img.shape[0], self.test_size[1] / img.shape[1])
img_info["ratio"] = ratio
img, _ = self.preproc(img, None, self.test_size)
img = torch.from_numpy(img).unsqueeze(0)
img = img.float()
if self.device == "gpu":
img = img.cuda()
if self.fp16:
img = img.half() # to FP16
with torch.no_grad():
outputs = self.model(img)
outputs = postprocess(
outputs, self.num_classes, self.confthre,
self.nmsthre
)
return outputs, img_info
def visual(self, output, img_info):
ratio = img_info["ratio"]
img = img_info["raw_img"]
if output is None:
return img
output = output.cpu()
bboxes = output[:, 0:4]
# preprocessing: resize
bboxes /= ratio
cls = output[:, 6]
scores = output[:, 4] * output[:, 5]
vis_res = vis(img, bboxes, scores, cls, self.confthre, self.cls_names)
return vis_res
|