File size: 5,629 Bytes
a292bf5 8f28ed6 a292bf5 f89be47 a292bf5 a757aaf 7fa23ea a292bf5 7fa23ea a292bf5 c5c0193 a292bf5 ac5ef0a 0bc3513 ac5ef0a a292bf5 7fa23ea a292bf5 7fa23ea a292bf5 4fd8a7f 04e0ad9 e891261 7fa23ea e891261 7fa23ea e891261 a292bf5 c4f0e59 a292bf5 |
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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import time
import os
import json
from loguru import logger
import cv2
import torch
from yolox.data.data_augment import ValTransform
from yolox.data.datasets import COCO_CLASSES
from yolox.exp import get_exp
from yolox.utils import fuse_model, get_model_info, postprocess, vis
IMAGE_EXT = [".jpg", ".jpeg", ".webp", ".bmp", ".png"]
def get_image_list(path):
image_names = []
for maindir, subdir, file_name_list in os.walk(path):
for filename in file_name_list:
apath = os.path.join(maindir, filename)
ext = os.path.splitext(apath)[1]
if ext in IMAGE_EXT:
image_names.append(apath)
return image_names
class Predictor(object):
def __init__(
self,
model,
exp,
cls_names=COCO_CLASSES,
trt_file=None,
decoder=None,
device="cpu",
fp16=False,
legacy=False,
):
self.model = model
self.cls_names = cls_names
self.decoder = decoder
self.num_classes = exp.num_classes
self.confthre = exp.test_conf
self.nmsthre = exp.nmsthre
self.test_size = exp.test_size
self.device = device
self.fp16 = fp16
self.preproc = ValTransform(legacy=legacy)
if trt_file is not None:
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load(trt_file))
x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()
self.model(x)
self.model = model_trt
def inference(self, img):
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
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])
pad_h =0 # ( self.test_size[0] - img.shape[0] * ratio ) / 2
pad_w =0 # ( self.test_size[1] - img.shape[1] * ratio ) / 2
img_info["pad"] = (pad_w, pad_h)
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():
t0 = time.time()
outputs = self.model(img)
if self.decoder is not None:
outputs = self.decoder(outputs, dtype=outputs.type())
outputs = postprocess(
outputs, self.num_classes, self.confthre,
self.nmsthre, class_agnostic=True
)
logger.info("Infer time: {:.4f}s".format(time.time() - t0))
return outputs, img_info
def build_predictor(
exp_file, model_name, ckpt_path, device="cpu", fp16=False, fuse=False, trt=False, conf=0, nms=0, tsize=None
):
# load experiment
exp = get_exp(exp_file, model_name)
if conf is not None:
exp.test_conf = conf
if nms is not None:
exp.nmsthre = nms
if tsize is not None:
exp.test_size = (tsize, tsize)
# create & initialize model
model = exp.get_model()
if device == "gpu":
model.cuda()
if fp16:
model.half()
model.eval()
logger.info("Model Summary: {}".format(get_model_info(model, exp.test_size)))
logger.info("loading checkpoint")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
# load the model state dict
model.load_state_dict(ckpt["model"])
logger.info("loaded checkpoint done.")
predictor = Predictor(
model, exp, COCO_CLASSES,
None, decoder=None,
device=device, fp16=fp16, legacy=False
)
return predictor
def run_detection(predictor, path):
# COCO output format: { images: [{id: 0, filename: "x.jpg"}, ...],
# annotations: [{id: 0, image_id: 0, bbox: [0 0 0 0], score: 0.35, class: 1}, ... ] }
if os.path.isdir(path):
files = get_image_list(path)
else:
files = [path]
files.sort()
img_list = []
ann_list = []
for img_id, image_name in enumerate(files):
outputs, img_info = predictor.inference(image_name)
ratio = img_info["ratio"]
pad_w, pad_h = img_info["pad"]
img_entry = {"id": img_id,
"filename": image_name }
img_list.append(img_entry)
if outputs[0] is not None:
for id, output in enumerate(outputs[0]):
print(output)
x1, y1, x2, y2 = output[:4]
x1 = (x1 - pad_w) / ratio
y1 = (y1 - pad_h) / ratio
x2 = (x2 - pad_w) / ratio
y2 = (y2 - pad_h) / ratio
ann_entry = {"id": id,
"image_id": img_id,
"bbox": [float(x1), float(y1), float(x2), float(y2)],
"cls": output[6].item(),
"score": (output[4] * output[5]).item() }
ann_list.append(ann_entry)
data_dict = { "images": img_list,
"annotations": ann_list
}
with open(f"{path}/results.json", 'w') as f:
json.dump(data_dict, f)
return f"{path}/results.json"
|