1717 lines
75 KiB
Python
Executable File
1717 lines
75 KiB
Python
Executable File
"""
|
|
测试代码：
|
|
1.-dw 输入目标检测权重，为pt文件，必要选项
|
|
2.-cw 输入细粒度识别权重，为onnx文件，必要选项
|
|
3.-c 输入细粒度识别名录，为txt文件，中文名字，每个中文用引号表示，中间用“，”隔开
|
|
4.-m 系统执行模式，分两种，一种为推理图片，直接生成结果“inf”，一种为AI label，生成xml文件，“ai”
|
|
5.-s 源文件路径，可以是文件夹也可以是图片文件，如果是AI label，则在同级目录下面生成一个xml文件。，可以一直输入单张图片地址进行推理“always”，则源文件路径参数失效，需要输入1即可
|
|
6.-o 存储根路径
|
|
CUDA_VISIBLE_DEVICES=7 ./NewDetAiLabelFine -dw /data5/Train_Models/Detection_weigth/3090/202405/20240501-17413-pixel50-GHY/weights/best.pt -cw /data5/Train_Models/mmcls_onnx/CMB_GHY_Animal_52_1/convnext-base_4xb32_cmb200/end2end.onnx -c /data5/CAS_Project/CAS_ProjectAnimal/CMB_ProjectAnimal/CMB_GHY_Animal_52/animalListCNlist.txt -m "ai" -s /data5/Temp_ProjectData/mdh/GHY/20240604/NET_camera_images_01/ -o /data5/Hatcher/data/detection/testresult5
|
|
"""
|
|
import imp
|
|
from unicodedata import name
|
|
import torch,os
|
|
import cv2,time
|
|
import warnings
|
|
from PIL import Image
|
|
import onnxruntime as ort
|
|
import numpy as np
|
|
import torchvision
|
|
import argparse,mimetypes
|
|
from tqdm import tqdm
|
|
from scipy.io import loadmat
|
|
import cv2
|
|
import json
|
|
from copy import deepcopy
|
|
from torch import Tensor
|
|
from configs.configs import model_head
|
|
from tqdm import tqdm
|
|
import os, sys, shutil
|
|
import torch
|
|
import platform,subprocess
|
|
import math,urllib
|
|
import random,requests
|
|
from pathlib import Path
|
|
from torch import nn
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from vizer.draw import draw_boxes
|
|
from concurrent.futures import ThreadPoolExecutor,as_completed
|
|
try:
|
|
import thop # for FLOPs computation
|
|
except ImportError:
|
|
thop = None
|
|
# from ssd.config import cfg
|
|
# from ssd.data.datasets import VOCDataset
|
|
import numpy as np
|
|
import pkg_resources as pkg
|
|
# from ssd.data.transforms import build_transforms
|
|
# from ssd.modeling.detector import build_detection_model
|
|
# from ssd.utils.checkpoint import CheckPointer
|
|
#from fine.models.LoadModel import MainModel
|
|
from fine.fineConfig import LoadConfig, load_data_transformers
|
|
from torch.autograd import Variable
|
|
from creat_xml_fordetection import make_xml
|
|
import threading
|
|
from common import global_para as gp
|
|
import time, statistics
|
|
from typing import Tuple
|
|
import torch.nn.functional as F
|
|
# root_path = os.getcwd()
|
|
# sys.path.insert(0,root_path+"/Y5")
|
|
|
|
#from Y5_old.utils.plots import plot_one_box
|
|
root_path = os.getcwd()
|
|
sys.path.insert(0,root_path+"/Y5")
|
|
import matplotlib.pyplot as plt
|
|
import Y5.models
|
|
DebugPrint_flag = True
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description='Hatcher detection inference ')
|
|
parser.add_argument('-dw', '--detectionweights', type=str,default="/data5/PartDataset/007_BJNHZ/0018_20240608/20231211/best.pt", help='Input weights (.pt) file path (required)')
|
|
parser.add_argument('-cw', '--fineweights', type=str,default="/data5/PartDataset/007_BJNHZ/0018_20240608/20231211/convnext_base_in22k-CMB_BJNHZ_20230202_31_84590-RP-RC-RP.onnx", help='Input weights (.onnx) file path (required)')
|
|
parser.add_argument('-c', '--classes',type=str,default="/data5/PartDataset/007_BJNHZ/0018_20240608/20231211/CMB_BJNHZ_20230911_31.txt" , help='Input fine class label')
|
|
parser.add_argument('-m','--mode', default="ai",choices=['inf', 'ai'], help='inference mode')
|
|
parser.add_argument('-s', '--source', type=str,default="/data5/Hatcher/data/detection/test1/", help='Input test images file path (required)')
|
|
parser.add_argument(
|
|
'-o', '--out', type=str, default="/data5/Hatcher/data/detection/testresult", help='save path')
|
|
args = parser.parse_args()
|
|
if not os.path.isfile(args.detectionweights):
|
|
raise SystemExit('Invalid weights file')
|
|
return args.detectionweights, args.fineweights,args.classes,args.mode, args.source,args.out
|
|
|
|
######
|
|
def cv2AddChineseText(draw, text ,position, textColor=(0, 255, 0), textSize=30):
|
|
# tp1 = time.time()
|
|
# if (isinstance(img, np.ndarray)): # ?D??¨º?¡¤?OpenCV¨ª???¨¤¨¤D¨ª
|
|
# img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
|
# tp2 = time.time()
|
|
# if gp.DebugPrint_flag:
|
|
# print("tp1 is ", (tp2 - tp1) * 1000)
|
|
# draw = ImageDraw.Draw(img)
|
|
tp2 = time.time()
|
|
fontStyle = ImageFont.truetype(
|
|
"./Font/msyh.ttc", textSize, encoding="utf-8")
|
|
tp3 = time.time()
|
|
if gp.DebugPrint_flag:
|
|
print("tp2 is ", (tp3 - tp2) * 1000)
|
|
|
|
draw.text(position, text, textColor, font=fontStyle)
|
|
|
|
tp4 = time.time()
|
|
if gp.DebugPrint_flag:
|
|
print("tp3 is ", (tp4 - tp3) * 1000)
|
|
return draw#cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
|
|
def plot_one_box(x, draw, height,width,color=None, label=None, Index = None, line_thickness=3):
|
|
# Plots one bounding box on image img
|
|
tl = line_thickness or round(0.002 * (height + width) / 2) + 1 # line/font thickness
|
|
color = color or tuple(random.randint(0, 255) for _ in range(3))
|
|
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
|
#ÏÔʾindex
|
|
|
|
|
|
# draw.text((c1[0]-20, c1[1]), str(Index), (0, 0, 255),font=fontStyle)
|
|
draw.rectangle((c1, c2), fill=None, outline=color, width=tl)
|
|
if label:
|
|
tf = max(tl - 1, 1) # font thickness
|
|
t_size = cv2.getTextSize(label, fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=gp.fontScale , thickness=tl)[0]
|
|
print(label,t_size[0],t_size[1])
|
|
c1 = c1[0],c1[1]-t_size[1]
|
|
c2 = c1[0] + t_size[0], c1[1]+t_size[1]
|
|
draw.rectangle((c1, c2), fill=color, outline=color, width=tl)
|
|
#cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
|
|
#cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
|
#if not label == None:
|
|
|
|
draw = cv2AddChineseText(draw, label,(c1[0], c1[1]), textColor = (225, 255, 255), textSize=tf*17)
|
|
#img = cv2AddChineseText(img, "ÊýÁ¿: {}".format(str("011")), (2000,100), textColor=(0, 255, 0), textSize=80)
|
|
return draw
|
|
def clip_coords(boxes, shape):
|
|
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
|
if isinstance(boxes, torch.Tensor): # faster individually
|
|
boxes[:, 0].clamp_(0, shape[1]) # x1
|
|
boxes[:, 1].clamp_(0, shape[0]) # y1
|
|
boxes[:, 2].clamp_(0, shape[1]) # x2
|
|
boxes[:, 3].clamp_(0, shape[0]) # y2
|
|
else: # np.array (faster grouped)
|
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
|
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
|
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
|
# Rescale coords (xyxy) from img1_shape to img0_shape
|
|
if ratio_pad is None: # calculate from img0_shape
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
|
else:
|
|
gain = ratio_pad[0][0]
|
|
pad = ratio_pad[1]
|
|
|
|
coords[:, [0, 2]] -= pad[0] # x padding
|
|
coords[:, [1, 3]] -= pad[1] # y padding
|
|
coords[:, :4] /= gain
|
|
clip_coords(coords, img0_shape)
|
|
return coords
|
|
def box_area(box):
|
|
# box = xyxy(4,n)
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
def _upcast(t: Tensor) -> Tensor:
|
|
# Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
|
|
if t.is_floating_point():
|
|
return t if t.dtype in (torch.float32, torch.float64) else t.float()
|
|
else:
|
|
return t if t.dtype in (torch.int32, torch.int64) else t.int()
|
|
def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]:
|
|
area1 = box_area(boxes1)
|
|
area2 = box_area(boxes2)
|
|
|
|
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
|
|
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
|
|
|
|
wh = _upcast(rb - lt).clamp(min=0) # [N,M,2]
|
|
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
|
|
|
|
union = area1[:, None] + area2 - inter
|
|
|
|
return inter, union
|
|
def xywh2xyxy(x):
|
|
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
|
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
|
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
|
|
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
|
|
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
|
|
return y
|
|
def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
|
|
"""
|
|
Return intersection-over-union (Jaccard index) of boxes.
|
|
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
|
|
Args:
|
|
boxes1 (Tensor[N, 4])
|
|
boxes2 (Tensor[M, 4])
|
|
|
|
Returns:
|
|
iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
|
|
"""
|
|
inter, union = _box_inter_union(boxes1, boxes2)
|
|
iou = inter / union
|
|
return iou
|
|
def non_max_suppression(prediction,
|
|
conf_thres=0.25,
|
|
iou_thres=0.45,
|
|
classes=None,
|
|
agnostic=False,
|
|
multi_label=False,
|
|
labels=(),
|
|
max_det=300):
|
|
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes
|
|
|
|
Returns:
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
"""
|
|
|
|
bs = prediction.shape[0] # batch size
|
|
nc = prediction.shape[2] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Checks
|
|
assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
|
|
assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
|
|
|
|
# Settings
|
|
# min_wh = 2 # (pixels) minimum box width and height
|
|
max_wh = 7680 # (pixels) maximum box width and height
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
|
time_limit = 0.3 + 0.03 * bs # seconds to quit after
|
|
redundant = True # require redundant detections
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
merge = False # use merge-NMS
|
|
|
|
t = time.time()
|
|
output = [torch.zeros((0, 6), device=prediction.device)] * bs
|
|
for xi, x in enumerate(prediction): # image index, image inference
|
|
# Apply constraints
|
|
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
|
x = x[xc[xi]] # confidence
|
|
|
|
# Cat apriori labels if autolabelling
|
|
if labels and len(labels[xi]):
|
|
lb = labels[xi]
|
|
v = torch.zeros((len(lb), nc + 5), device=x.device)
|
|
v[:, :4] = lb[:, 1:5] # box
|
|
v[:, 4] = 1.0 # conf
|
|
v[range(len(lb)), lb[:, 0].long() + 5] = 1.0 # cls
|
|
x = torch.cat((x, v), 0)
|
|
|
|
# If none remain process next image
|
|
if not x.shape[0]:
|
|
continue
|
|
|
|
# Compute conf
|
|
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
|
|
|
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
|
box = xywh2xyxy(x[:, :4])
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
if multi_label:
|
|
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
|
else: # best class only
|
|
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
|
|
|
# Filter by class
|
|
if classes is not None:
|
|
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
|
|
|
# Apply finite constraint
|
|
# if not torch.isfinite(x).all():
|
|
# x = x[torch.isfinite(x).all(1)]
|
|
|
|
# Check shape
|
|
n = x.shape[0] # number of boxes
|
|
if not n: # no boxes
|
|
continue
|
|
elif n > max_nms: # excess boxes
|
|
x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
|
|
|
# Batched NMS
|
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
|
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
|
weights = iou * scores[None] # box weights
|
|
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
|
if redundant:
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
|
|
output[xi] = x[i]
|
|
if (time.time() - t) > time_limit:
|
|
|
|
break # time limit exceeded
|
|
|
|
return output
|
|
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
|
# Resize and pad image while meeting stride-multiple constraints
|
|
shape = im.shape[:2] # current shape [height, width]
|
|
if isinstance(new_shape, int):
|
|
new_shape = (new_shape, new_shape)
|
|
|
|
# Scale ratio (new / old)
|
|
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
|
if not scaleup: # only scale down, do not scale up (for better val mAP)
|
|
r = min(r, 1.0)
|
|
|
|
# Compute padding
|
|
ratio = r, r # width, height ratios
|
|
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
|
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
|
# if auto: # minimum rectangle
|
|
# dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
|
# elif scaleFill: # stretch
|
|
# dw, dh = 0.0, 0.0
|
|
# new_unpad = (new_shape[1], new_shape[0])
|
|
# ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
|
|
|
|
dw /= 2 # divide padding into 2 sides
|
|
dh /= 2
|
|
|
|
if shape[::-1] != new_unpad: # resize
|
|
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
|
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
|
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
|
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
|
return im, ratio, (dw, dh)
|
|
def check_anchor_order(m):
|
|
# Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
|
|
a = m.anchors.prod(-1).mean(-1).view(-1) # mean anchor area per output layer
|
|
da = a[-1] - a[0] # delta a
|
|
ds = m.stride[-1] - m.stride[0] # delta s
|
|
if da and (da.sign() != ds.sign()): # same order
|
|
m.anchors[:] = m.anchors.flip(0)
|
|
def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False, verbose=False):
|
|
# Check version vs. required version
|
|
current, minimum = (pkg.parse_version(x) for x in (current, minimum))
|
|
result = (current == minimum) if pinned else (current >= minimum) # bool
|
|
s = f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed' # string
|
|
if hard:
|
|
assert result, s # assert min requirements met
|
|
if verbose and not result:
|
|
pass
|
|
return result
|
|
class Detect(nn.Module):
|
|
stride = None # strides computed during build
|
|
onnx_dynamic = False # ONNX export parameter
|
|
export = False # export mode
|
|
|
|
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
|
|
super().__init__()
|
|
self.nc = nc # number of classes
|
|
self.no = nc + 5 # number of outputs per anchor
|
|
self.nl = len(anchors) # number of detection layers
|
|
self.na = len(anchors[0]) // 2 # number of anchors
|
|
self.grid = [torch.zeros(1)] * self.nl # init grid
|
|
self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
|
|
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
|
|
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
|
|
self.inplace = inplace # use in-place ops (e.g. slice assignment)
|
|
|
|
def forward(self, x):
|
|
z = [] # inference output
|
|
for i in range(self.nl):
|
|
x[i] = self.m[i](x[i]) # conv
|
|
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
|
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
|
|
|
if not self.training: # inference
|
|
if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
|
|
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
|
|
|
|
y = x[i].sigmoid()
|
|
if self.inplace:
|
|
y[..., 0:2] = (y[..., 0:2] * 2 + self.grid[i]) * self.stride[i] # xy
|
|
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
|
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
|
|
xy, wh, conf = y.split((2, 2, self.nc + 1), 4) # y.tensor_split((2, 4, 5), 4) # torch 1.8.0
|
|
xy = (xy * 2 + self.grid[i]) * self.stride[i] # xy
|
|
wh = (wh * 2) ** 2 * self.anchor_grid[i] # wh
|
|
y = torch.cat((xy, wh, conf), 4)
|
|
z.append(y.view(bs, -1, self.no))
|
|
|
|
return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
|
|
|
|
def _make_grid(self, nx=20, ny=20, i=0):
|
|
d = self.anchors[i].device
|
|
t = self.anchors[i].dtype
|
|
shape = 1, self.na, ny, nx, 2 # grid shape
|
|
y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
|
|
if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
|
|
yv, xv = torch.meshgrid(y, x, indexing='ij')
|
|
else:
|
|
yv, xv = torch.meshgrid(y, x)
|
|
grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 # add grid offset, i.e. y = 2.0 * x - 0.5
|
|
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
|
return grid, anchor_grid
|
|
|
|
def initialize_weights(model):
|
|
for m in model.modules():
|
|
t = type(m)
|
|
if t is nn.Conv2d:
|
|
pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
|
elif t is nn.BatchNorm2d:
|
|
m.eps = 1e-3
|
|
m.momentum = 0.03
|
|
elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
|
|
m.inplace = True
|
|
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
|
# Scales img(bs,3,y,x) by ratio constrained to gs-multiple
|
|
if ratio == 1.0:
|
|
return img
|
|
h, w = img.shape[2:]
|
|
s = (int(h * ratio), int(w * ratio)) # new size
|
|
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
|
if not same_shape: # pad/crop img
|
|
h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
|
|
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
|
def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')):
|
|
"""
|
|
x: Features to be visualized
|
|
module_type: Module type
|
|
stage: Module stage within model
|
|
n: Maximum number of feature maps to plot
|
|
save_dir: Directory to save results
|
|
"""
|
|
if 'Detect' not in module_type:
|
|
batch, channels, height, width = x.shape # batch, channels, height, width
|
|
if height > 1 and width > 1:
|
|
f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
|
|
|
|
blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
|
|
n = min(n, channels) # number of plots
|
|
fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols
|
|
ax = ax.ravel()
|
|
plt.subplots_adjust(wspace=0.05, hspace=0.05)
|
|
for i in range(n):
|
|
ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
|
|
ax[i].axis('off')
|
|
|
|
|
|
plt.savefig(f, dpi=300, bbox_inches='tight')
|
|
plt.close()
|
|
np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save
|
|
def time_sync():
|
|
# PyTorch-accurate time
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
return time.time()
|
|
def autopad(k, p=None): # kernel, padding
|
|
# Pad to 'same'
|
|
if p is None:
|
|
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
|
|
return p
|
|
class Conv(nn.Module):
|
|
# Standard convolution
|
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
|
|
self.bn = nn.BatchNorm2d(c2)
|
|
self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
|
|
|
|
def forward(self, x):
|
|
return self.act(self.bn(self.conv(x)))
|
|
|
|
def forward_fuse(self, x):
|
|
return self.act(self.conv(x))
|
|
|
|
|
|
class DWConv(Conv):
|
|
# Depth-wise convolution class
|
|
def __init__(self, c1, c2, k=1, s=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
|
super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
|
|
|
|
|
|
class DWConvTranspose2d(nn.ConvTranspose2d):
|
|
# Depth-wise transpose convolution class
|
|
def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0): # ch_in, ch_out, kernel, stride, padding, padding_out
|
|
super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
|
|
|
|
|
|
class TransformerLayer(nn.Module):
|
|
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
|
|
def __init__(self, c, num_heads):
|
|
super().__init__()
|
|
self.q = nn.Linear(c, c, bias=False)
|
|
self.k = nn.Linear(c, c, bias=False)
|
|
self.v = nn.Linear(c, c, bias=False)
|
|
self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
|
|
self.fc1 = nn.Linear(c, c, bias=False)
|
|
self.fc2 = nn.Linear(c, c, bias=False)
|
|
|
|
def forward(self, x):
|
|
x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
|
|
x = self.fc2(self.fc1(x)) + x
|
|
return x
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
# Vision Transformer https://arxiv.org/abs/2010.11929
|
|
def __init__(self, c1, c2, num_heads, num_layers):
|
|
super().__init__()
|
|
self.conv = None
|
|
if c1 != c2:
|
|
self.conv = Conv(c1, c2)
|
|
self.linear = nn.Linear(c2, c2) # learnable position embedding
|
|
self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
|
|
self.c2 = c2
|
|
|
|
def forward(self, x):
|
|
if self.conv is not None:
|
|
x = self.conv(x)
|
|
b, _, w, h = x.shape
|
|
p = x.flatten(2).permute(2, 0, 1)
|
|
return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
|
|
|
|
|
|
class Bottleneck(nn.Module):
|
|
# Standard bottleneck
|
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
|
|
super().__init__()
|
|
c_ = int(c2 * e) # hidden channels
|
|
self.cv1 = Conv(c1, c_, 1, 1)
|
|
self.cv2 = Conv(c_, c2, 3, 1, g=g)
|
|
self.add = shortcut and c1 == c2
|
|
|
|
def forward(self, x):
|
|
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
|
|
|
|
|
|
class BottleneckCSP(nn.Module):
|
|
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
|
super().__init__()
|
|
c_ = int(c2 * e) # hidden channels
|
|
self.cv1 = Conv(c1, c_, 1, 1)
|
|
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
|
|
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
|
|
self.cv4 = Conv(2 * c_, c2, 1, 1)
|
|
self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
|
|
self.act = nn.SiLU()
|
|
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
|
|
|
|
def forward(self, x):
|
|
y1 = self.cv3(self.m(self.cv1(x)))
|
|
y2 = self.cv2(x)
|
|
return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))
|
|
|
|
|
|
class CrossConv(nn.Module):
|
|
# Cross Convolution Downsample
|
|
def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
|
|
# ch_in, ch_out, kernel, stride, groups, expansion, shortcut
|
|
super().__init__()
|
|
c_ = int(c2 * e) # hidden channels
|
|
self.cv1 = Conv(c1, c_, (1, k), (1, s))
|
|
self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
|
|
self.add = shortcut and c1 == c2
|
|
|
|
def forward(self, x):
|
|
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
|
|
|
|
|
|
class C3(nn.Module):
|
|
# CSP Bottleneck with 3 convolutions
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
|
super().__init__()
|
|
c_ = int(c2 * e) # hidden channels
|
|
self.cv1 = Conv(c1, c_, 1, 1)
|
|
self.cv2 = Conv(c1, c_, 1, 1)
|
|
self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
|
|
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
|
|
|
|
def forward(self, x):
|
|
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
|
|
|
|
|
|
class C3x(C3):
|
|
# C3 module with cross-convolutions
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
c_ = int(c2 * e)
|
|
self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
|
|
|
|
|
|
class C3TR(C3):
|
|
# C3 module with TransformerBlock()
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
c_ = int(c2 * e)
|
|
self.m = TransformerBlock(c_, c_, 4, n)
|
|
|
|
|
|
class C3SPP(C3):
|
|
# C3 module with SPP()
|
|
def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
c_ = int(c2 * e)
|
|
self.m = SPP(c_, c_, k)
|
|
|
|
|
|
class C3Ghost(C3):
|
|
# C3 module with GhostBottleneck()
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
c_ = int(c2 * e) # hidden channels
|
|
self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
|
|
|
|
|
|
class SPP(nn.Module):
|
|
# Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
|
|
def __init__(self, c1, c2, k=(5, 9, 13)):
|
|
super().__init__()
|
|
c_ = c1 // 2 # hidden channels
|
|
self.cv1 = Conv(c1, c_, 1, 1)
|
|
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
|
|
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
|
|
|
|
def forward(self, x):
|
|
x = self.cv1(x)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
|
|
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
|
|
|
|
FONT = 'Arial.ttf'
|
|
class SPPF(nn.Module):
|
|
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
|
|
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
|
|
super().__init__()
|
|
c_ = c1 // 2 # hidden channels
|
|
self.cv1 = Conv(c1, c_, 1, 1)
|
|
self.cv2 = Conv(c_ * 4, c2, 1, 1)
|
|
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
|
|
|
|
def forward(self, x):
|
|
x = self.cv1(x)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
|
|
y1 = self.m(x)
|
|
y2 = self.m(y1)
|
|
return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
|
|
|
|
|
|
class Focus(nn.Module):
|
|
# Focus wh information into c-space
|
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
|
super().__init__()
|
|
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
|
|
# self.contract = Contract(gain=2)
|
|
|
|
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
|
|
return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
|
|
# return self.conv(self.contract(x))
|
|
|
|
|
|
class GhostConv(nn.Module):
|
|
# Ghost Convolution https://github.com/huawei-noah/ghostnet
|
|
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
|
super().__init__()
|
|
c_ = c2 // 2 # hidden channels
|
|
self.cv1 = Conv(c1, c_, k, s, None, g, act)
|
|
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
|
|
|
|
def forward(self, x):
|
|
y = self.cv1(x)
|
|
return torch.cat((y, self.cv2(y)), 1)
|
|
|
|
class MixConv2d(nn.Module):
|
|
# Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595
|
|
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kernel, stride, ch_strategy
|
|
super().__init__()
|
|
n = len(k) # number of convolutions
|
|
if equal_ch: # equal c_ per group
|
|
i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices
|
|
c_ = [(i == g).sum() for g in range(n)] # intermediate channels
|
|
else: # equal weight.numel() per group
|
|
b = [c2] + [0] * n
|
|
a = np.eye(n + 1, n, k=-1)
|
|
a -= np.roll(a, 1, axis=1)
|
|
a *= np.array(k) ** 2
|
|
a[0] = 1
|
|
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
|
|
|
self.m = nn.ModuleList([
|
|
nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)])
|
|
self.bn = nn.BatchNorm2d(c2)
|
|
self.act = nn.SiLU()
|
|
|
|
def forward(self, x):
|
|
return self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
|
class GhostBottleneck(nn.Module):
|
|
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
|
|
def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
|
|
super().__init__()
|
|
c_ = c2 // 2
|
|
self.conv = nn.Sequential(
|
|
GhostConv(c1, c_, 1, 1), # pw
|
|
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
|
|
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
|
|
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
|
|
act=False)) if s == 2 else nn.Identity()
|
|
|
|
def forward(self, x):
|
|
return self.conv(x) + self.shortcut(x)
|
|
|
|
|
|
class Contract(nn.Module):
|
|
# Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
|
|
def __init__(self, gain=2):
|
|
super().__init__()
|
|
self.gain = gain
|
|
|
|
def forward(self, x):
|
|
b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
|
|
s = self.gain
|
|
x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)
|
|
x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
|
|
return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)
|
|
|
|
|
|
class Expand(nn.Module):
|
|
# Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
|
|
def __init__(self, gain=2):
|
|
super().__init__()
|
|
self.gain = gain
|
|
|
|
def forward(self, x):
|
|
b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
|
|
s = self.gain
|
|
x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80)
|
|
x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
|
|
return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160)
|
|
|
|
|
|
class Concat(nn.Module):
|
|
# Concatenate a list of tensors along dimension
|
|
def __init__(self, dimension=1):
|
|
super().__init__()
|
|
self.d = dimension
|
|
|
|
def forward(self, x):
|
|
return torch.cat(x, self.d)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def xyxy2xywh(x):
|
|
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
|
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
|
y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
|
|
y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
|
|
y[:, 2] = x[:, 2] - x[:, 0] # width
|
|
y[:, 3] = x[:, 3] - x[:, 1] # height
|
|
return y
|
|
def is_ascii(s=''):
|
|
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
|
|
s = str(s) # convert list, tuple, None, etc. to str
|
|
return len(s.encode().decode('ascii', 'ignore')) == len(s)
|
|
def is_writeable(dir, test=False):
|
|
# Return True if directory has write permissions, test opening a file with write permissions if test=True
|
|
if not test:
|
|
return os.access(dir, os.R_OK) # possible issues on Windows
|
|
file = Path(dir) / 'tmp.txt'
|
|
try:
|
|
with open(file, 'w'): # open file with write permissions
|
|
pass
|
|
file.unlink() # remove file
|
|
return True
|
|
except OSError:
|
|
return False
|
|
def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
|
|
# Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
|
|
env = os.getenv(env_var)
|
|
if env:
|
|
path = Path(env) # use environment variable
|
|
else:
|
|
cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs
|
|
path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir
|
|
path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable
|
|
path.mkdir(exist_ok=True) # make if required
|
|
return path
|
|
CONFIG_DIR = user_config_dir()
|
|
def check_font(font=FONT, progress=False):
|
|
# Download font to CONFIG_DIR if necessary
|
|
font = Path(font)
|
|
file = CONFIG_DIR / font.name
|
|
if not font.exists() and not file.exists():
|
|
url = "https://ultralytics.com/assets/" + font.name
|
|
|
|
torch.hub.download_url_to_file(url, str(file), progress=progress)
|
|
|
|
def check_pil_font(font=FONT, size=10):
|
|
# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
|
|
font = Path(font)
|
|
font = font if font.exists() else (CONFIG_DIR / font.name)
|
|
try:
|
|
return ImageFont.truetype(str(font) if font.exists() else font.name, size)
|
|
except Exception: # download if missing
|
|
try:
|
|
check_font(font)
|
|
return ImageFont.truetype(str(font), size)
|
|
except TypeError:
|
|
pass # known issue https://github.com/ultralytics/yolov5/issues/5374
|
|
|
|
|
|
class Annotator:
|
|
# YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
|
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
|
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
|
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
|
self.pil = pil or non_ascii
|
|
if self.pil: # use PIL
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
self.draw = ImageDraw.Draw(self.im)
|
|
self.font = check_pil_font(font='Arial.Unicode.ttf' if non_ascii else font,
|
|
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
|
else: # use cv2
|
|
self.im = im
|
|
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
|
|
|
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
|
# Add one xyxy box to image with label
|
|
if self.pil or not is_ascii(label):
|
|
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
|
if label:
|
|
w, h = self.font.getsize(label) # text width, height
|
|
outside = box[1] - h >= 0 # label fits outside box
|
|
self.draw.rectangle(
|
|
(box[0], box[1] - h if outside else box[1], box[0] + w + 1,
|
|
box[1] + 1 if outside else box[1] + h + 1),
|
|
fill=color,
|
|
)
|
|
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
|
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
|
else: # cv2
|
|
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
|
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
|
|
if label:
|
|
tf = max(self.lw - 1, 1) # font thickness
|
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
|
outside = p1[1] - h >= 3
|
|
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
|
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
|
cv2.putText(self.im,
|
|
label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
|
|
0,
|
|
self.lw / 3,
|
|
txt_color,
|
|
thickness=tf,
|
|
lineType=cv2.LINE_AA)
|
|
|
|
def rectangle(self, xy, fill=None, outline=None, width=1):
|
|
# Add rectangle to image (PIL-only)
|
|
self.draw.rectangle(xy, fill, outline, width)
|
|
|
|
def text(self, xy, text, txt_color=(255, 255, 255)):
|
|
# Add text to image (PIL-only)
|
|
w, h = self.font.getsize(text) # text width, height
|
|
self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)
|
|
|
|
def result(self):
|
|
# Return annotated image as array
|
|
return np.asarray(self.im)
|
|
|
|
class Detections:
|
|
# YOLOv5 detections class for inference results
|
|
def __init__(self, imgs, pred, files, times=(0, 0, 0, 0), names=None, shape=None):
|
|
super().__init__()
|
|
d = pred[0].device # device
|
|
gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in imgs] # normalizations
|
|
self.imgs = imgs # list of images as numpy arrays
|
|
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
|
|
self.names = names # class names
|
|
self.files = files # image filenames
|
|
self.times = times # profiling times
|
|
self.xyxy = pred # xyxy pixels
|
|
self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
|
|
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
|
|
self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
|
|
self.n = len(self.pred) # number of images (batch size)
|
|
self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
|
|
self.s = shape # inference BCHW shape
|
|
|
|
# def display(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
|
|
# crops = []
|
|
# for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
|
# s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
|
|
# if pred.shape[0]:
|
|
# for c in pred[:, -1].unique():
|
|
# n = (pred[:, -1] == c).sum() # detections per class
|
|
# s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
|
# if show or save or render or crop:
|
|
# annotator = Annotator(im, example=str(self.names))
|
|
# for *box, conf, cls in reversed(pred): # xyxy, confidence, class
|
|
# label = f'{self.names[int(cls)]} {conf:.2f}'
|
|
# if crop:
|
|
# file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
|
|
# crops.append({
|
|
# 'box': box,
|
|
# 'conf': conf,
|
|
# 'cls': cls,
|
|
# 'label': label,
|
|
# 'im': save_one_box(box, im, file=file, save=save)})
|
|
# else: # all others
|
|
# annotator.box_label(box, label if labels else '', color=colors(cls))
|
|
# im = annotator.im
|
|
# else:
|
|
# s += '(no detections)'
|
|
|
|
# im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
|
|
# if pprint:
|
|
# print(s.rstrip(', '))
|
|
# if show:
|
|
# im.show(self.files[i]) # show
|
|
# if save:
|
|
# f = self.files[i]
|
|
# im.save(save_dir / f) # save
|
|
# if i == self.n - 1:
|
|
# pass
|
|
# #LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
|
|
# if render:
|
|
# self.imgs[i] = np.asarray(im)
|
|
# if crop:
|
|
# if save:
|
|
# pass
|
|
# #LOGGER.info(f'Saved results to {save_dir}\n')
|
|
# return crops
|
|
|
|
# def print(self):
|
|
# self.display(pprint=True) # print results
|
|
# print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)
|
|
|
|
# def show(self, labels=True):
|
|
# self.display(show=True, labels=labels) # show results
|
|
|
|
|
|
|
|
|
|
# def render(self, labels=True):
|
|
# self.display(render=True, labels=labels) # render results
|
|
# return self.imgs
|
|
|
|
# def pandas(self):
|
|
# # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
|
|
# new = copy(self) # return copy
|
|
# ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
|
|
# cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
|
|
# for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
|
|
# a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
|
|
# setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
|
|
# return new
|
|
|
|
def tolist(self):
|
|
# return a list of Detections objects, i.e. 'for result in results.tolist():'
|
|
r = range(self.n) # iterable
|
|
x = [Detections([self.imgs[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
|
|
# for d in x:
|
|
# for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
|
|
# setattr(d, k, getattr(d, k)[0]) # pop out of list
|
|
return x
|
|
|
|
def __len__(self):
|
|
return self.n # override len(results)
|
|
|
|
def __str__(self):
|
|
self.print() # override print(results)
|
|
return ''
|
|
def fuse_conv_and_bn(conv, bn):
|
|
# Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
|
|
fusedconv = nn.Conv2d(conv.in_channels,
|
|
conv.out_channels,
|
|
kernel_size=conv.kernel_size,
|
|
stride=conv.stride,
|
|
padding=conv.padding,
|
|
groups=conv.groups,
|
|
bias=True).requires_grad_(False).to(conv.weight.device)
|
|
|
|
# Prepare filters
|
|
w_conv = conv.weight.clone().view(conv.out_channels, -1)
|
|
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
|
|
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
|
|
|
|
# Prepare spatial bias
|
|
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
|
|
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
|
|
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
|
|
|
|
return fusedconv
|
|
|
|
class Classify(nn.Module):
|
|
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
|
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
|
|
super().__init__()
|
|
self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
|
|
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
|
|
self.flat = nn.Flatten()
|
|
|
|
def forward(self, x):
|
|
z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
|
|
return self.flat(self.conv(z)) # flatten to x(b,c2)
|
|
class Model(nn.Module):
|
|
# YOLOv5 model
|
|
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
|
|
super().__init__()
|
|
if isinstance(cfg, dict):
|
|
self.yaml = cfg # model dict
|
|
else: # is *.yaml
|
|
import yaml # for torch hub
|
|
self.yaml_file = Path(cfg).name
|
|
with open(cfg, encoding='ascii', errors='ignore') as f:
|
|
self.yaml = yaml.safe_load(f) # model dict
|
|
|
|
# Define model
|
|
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
|
if nc and nc != self.yaml['nc']:
|
|
|
|
self.yaml['nc'] = nc # override yaml value
|
|
if anchors:
|
|
|
|
self.yaml['anchors'] = round(anchors) # override yaml value
|
|
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
|
|
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
|
|
self.inplace = self.yaml.get('inplace', True)
|
|
|
|
# Build strides, anchors
|
|
m = self.model[-1] # Detect()
|
|
if isinstance(m, Detect):
|
|
s = 256 # 2x min stride
|
|
m.inplace = self.inplace
|
|
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
|
|
check_anchor_order(m) # must be in pixel-space (not grid-space)
|
|
m.anchors /= m.stride.view(-1, 1, 1)
|
|
self.stride = m.stride
|
|
self._initialize_biases() # only run once
|
|
|
|
# Init weights, biases
|
|
initialize_weights(self)
|
|
self.info()
|
|
|
|
|
|
def forward(self, x, augment=False, profile=False, visualize=False):
|
|
if augment:
|
|
return self._forward_augment(x) # augmented inference, None
|
|
return self._forward_once(x, profile, visualize) # single-scale inference, train
|
|
|
|
def _forward_augment(self, x):
|
|
img_size = x.shape[-2:] # height, width
|
|
s = [1, 0.83, 0.67] # scales
|
|
f = [None, 3, None] # flips (2-ud, 3-lr)
|
|
y = [] # outputs
|
|
for si, fi in zip(s, f):
|
|
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
|
|
yi = self._forward_once(xi)[0] # forward
|
|
# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
|
|
yi = self._descale_pred(yi, fi, si, img_size)
|
|
y.append(yi)
|
|
y = self._clip_augmented(y) # clip augmented tails
|
|
return torch.cat(y, 1), None # augmented inference, train
|
|
|
|
def _forward_once(self, x, profile=False, visualize=False):
|
|
y, dt = [], [] # outputs
|
|
for m in self.model:
|
|
if m.f != -1: # if not from previous layer
|
|
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
|
|
if profile:
|
|
self._profile_one_layer(m, x, dt)
|
|
x = m(x) # run
|
|
y.append(x if m.i in self.save else None) # save output
|
|
if visualize:
|
|
feature_visualization(x, m.type, m.i, save_dir=visualize)
|
|
return x
|
|
|
|
def _descale_pred(self, p, flips, scale, img_size):
|
|
# de-scale predictions following augmented inference (inverse operation)
|
|
if self.inplace:
|
|
p[..., :4] /= scale # de-scale
|
|
if flips == 2:
|
|
p[..., 1] = img_size[0] - p[..., 1] # de-flip ud
|
|
elif flips == 3:
|
|
p[..., 0] = img_size[1] - p[..., 0] # de-flip lr
|
|
else:
|
|
x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # de-scale
|
|
if flips == 2:
|
|
y = img_size[0] - y # de-flip ud
|
|
elif flips == 3:
|
|
x = img_size[1] - x # de-flip lr
|
|
p = torch.cat((x, y, wh, p[..., 4:]), -1)
|
|
return p
|
|
|
|
def _clip_augmented(self, y):
|
|
# Clip YOLOv5 augmented inference tails
|
|
nl = self.model[-1].nl # number of detection layers (P3-P5)
|
|
g = sum(4 ** x for x in range(nl)) # grid points
|
|
e = 1 # exclude layer count
|
|
i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e)) # indices
|
|
y[0] = y[0][:, :-i] # large
|
|
i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices
|
|
y[-1] = y[-1][:, i:] # small
|
|
return y
|
|
|
|
def _profile_one_layer(self, m, x, dt):
|
|
c = isinstance(m, Detect) # is final layer, copy input as inplace fix
|
|
o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs
|
|
t = time_sync()
|
|
for _ in range(10):
|
|
m(x.copy() if c else x)
|
|
dt.append((time_sync() - t) * 100)
|
|
if m == self.model[0]:
|
|
pass
|
|
pass
|
|
if c:
|
|
pass
|
|
|
|
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
|
|
# https://arxiv.org/abs/1708.02002 section 3.3
|
|
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
|
m = self.model[-1] # Detect() module
|
|
for mi, s in zip(m.m, m.stride): # from
|
|
b = mi.bias.view(m.na, -1).detach() # conv.bias(255) to (3,85)
|
|
b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
|
b[:, 5:] += math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum()) # cls
|
|
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
|
|
|
def _print_biases(self):
|
|
m = self.model[-1] # Detect() module
|
|
for mi in m.m: # from
|
|
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
|
|
|
|
|
|
# def _print_weights(self):
|
|
# for m in self.model.modules():
|
|
# if type(m) is Bottleneck:
|
|
# LOGGER.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
|
|
|
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
|
|
|
for m in self.model.modules():
|
|
if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
|
|
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
|
delattr(m, 'bn') # remove batchnorm
|
|
m.forward = m.forward_fuse # update forward
|
|
self.info()
|
|
return self
|
|
|
|
def info(self, verbose=False, img_size=640): # print model information
|
|
pass
|
|
#model_info(self, verbose, img_size)
|
|
|
|
def _apply(self, fn):
|
|
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
|
|
self = super()._apply(fn)
|
|
m = self.model[-1] # Detect()
|
|
if isinstance(m, Detect):
|
|
m.stride = fn(m.stride)
|
|
m.grid = list(map(fn, m.grid))
|
|
if isinstance(m.anchor_grid, list):
|
|
m.anchor_grid = list(map(fn, m.anchor_grid))
|
|
return self
|
|
def make_divisible(x, divisor):
|
|
# Returns nearest x divisible by divisor
|
|
if isinstance(divisor, torch.Tensor):
|
|
divisor = int(divisor.max()) # to int
|
|
return math.ceil(x / divisor) * divisor
|
|
|
|
def parse_model(d, ch): # model_dict, input_channels(3)
|
|
#LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
|
|
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
|
|
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
|
|
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
|
|
|
|
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
|
|
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
|
|
m = eval(m) if isinstance(m, str) else m # eval strings
|
|
for j, a in enumerate(args):
|
|
try:
|
|
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
|
except NameError:
|
|
pass
|
|
|
|
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
|
|
if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
|
|
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x):
|
|
c1, c2 = ch[f], args[0]
|
|
if c2 != no: # if not output
|
|
c2 = make_divisible(c2 * gw, 8)
|
|
|
|
args = [c1, c2, *args[1:]]
|
|
if m in [BottleneckCSP, C3, C3TR, C3Ghost, C3x]:
|
|
args.insert(2, n) # number of repeats
|
|
n = 1
|
|
elif m is nn.BatchNorm2d:
|
|
args = [ch[f]]
|
|
elif m is Concat:
|
|
c2 = sum(ch[x] for x in f)
|
|
elif m is Detect:
|
|
args.append([ch[x] for x in f])
|
|
if isinstance(args[1], int): # number of anchors
|
|
args[1] = [list(range(args[1] * 2))] * len(f)
|
|
elif m is Contract:
|
|
c2 = ch[f] * args[0] ** 2
|
|
elif m is Expand:
|
|
c2 = ch[f] // args[0] ** 2
|
|
else:
|
|
c2 = ch[f]
|
|
|
|
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
|
|
t = str(m)[8:-2].replace('__main__.', '') # module type
|
|
np = sum(x.numel() for x in m_.parameters()) # number params
|
|
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
|
|
#LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}') # print
|
|
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
|
layers.append(m_)
|
|
if i == 0:
|
|
ch = []
|
|
ch.append(c2)
|
|
return nn.Sequential(*layers), sorted(save)
|
|
class Ensemble(nn.ModuleList):
|
|
# Ensemble of models
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x, augment=False, profile=False, visualize=False):
|
|
y = [module(x, augment, profile, visualize)[0] for module in self]
|
|
# y = torch.stack(y).max(0)[0] # max ensemble
|
|
# y = torch.stack(y).mean(0) # mean ensemble
|
|
y = torch.cat(y, 1) # nms ensemble
|
|
return y, None # inference, train output
|
|
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
|
|
# 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
|
|
#LOGGER.info(f'Downloading {url} to {file}...')
|
|
torch.hub.download_url_to_file(url, str(file), progress=True)
|
|
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
|
|
#LOGGER.info(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
|
|
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
|
|
#LOGGER.info(f"ERROR: {assert_msg}\n{error_msg}")
|
|
pass
|
|
#LOGGER.info('')
|
|
def attempt_download(file, repo='ultralytics/yolov5', release='v6.1'):
|
|
# Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v6.1', etc.
|
|
|
|
def github_assets(repository, version='latest'):
|
|
# Return GitHub repo tag (i.e. 'v6.1') and assets (i.e. ['yolov5s.pt', 'yolov5m.pt', ...])
|
|
if version != 'latest':
|
|
version = f'tags/{version}' # i.e. tags/v6.1
|
|
response = requests.get(f'https://api.github.com/repos/{repository}/releases/{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():
|
|
pass
|
|
#LOGGER.info(f'Found {url} locally at {file}') # file already exists
|
|
else:
|
|
safe_download(file=file, url=url, min_bytes=1E5)
|
|
return file
|
|
|
|
# GitHub assets
|
|
assets = [
|
|
'yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov5n6.pt', 'yolov5s6.pt',
|
|
'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.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:
|
|
url3 = 'https://drive.google.com/drive/folders/1EFQTEUeXWSFww0luse2jB9M1QNZQGwNl' # backup gdrive mirror
|
|
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} or {url3}')
|
|
|
|
return str(file)
|
|
def attempt_load(weights, device=None, inplace=True, fuse=True):
|
|
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
|
|
|
|
|
model = Ensemble()
|
|
for w in weights if isinstance(weights, list) else [weights]:
|
|
ckpt = torch.load(attempt_download(w), map_location='cpu') # load
|
|
ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model
|
|
model.append(ckpt.fuse().eval() if fuse else ckpt.eval()) # fused or un-fused model in eval mode
|
|
|
|
# Compatibility updates
|
|
for m in model.modules():
|
|
t = type(m)
|
|
if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model):
|
|
m.inplace = inplace # torch 1.7.0 compatibility
|
|
if t is Detect and not isinstance(m.anchor_grid, list):
|
|
delattr(m, 'anchor_grid')
|
|
setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl)
|
|
elif t is Conv:
|
|
m._non_persistent_buffers_set = set() # torch 1.6.0 compatibility
|
|
elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
|
|
m.recompute_scale_factor = None # torch 1.11.0 compatibility
|
|
|
|
if len(model) == 1:
|
|
return model[-1] # return model
|
|
print(f'Ensemble created with {weights}\n')
|
|
for k in 'names', 'nc', 'yaml':
|
|
setattr(model, k, getattr(model[0], k))
|
|
model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride
|
|
assert all(model[0].nc == m.nc for m in model), f'Models have different class counts: {[m.nc for m in model]}'
|
|
return model # return ensemble
|
|
|
|
|
|
######
|
|
|
|
|
|
def identify_path_type(path):
|
|
# 检查路径是否存在
|
|
if path =="always":
|
|
return "always"
|
|
|
|
if not os.path.exists(path):
|
|
return 0
|
|
|
|
# 检查是否为文件夹
|
|
if os.path.isdir(path):
|
|
items = os.listdir(path)
|
|
for item in items:
|
|
if os.path.isdir(os.path.join(path, item)):
|
|
return 4 # Found a subfolder
|
|
return 1
|
|
if os.path.splitext(path)[1].lower() == '.txt':
|
|
return 6
|
|
|
|
|
|
# 检查是否为文件
|
|
if os.path.isfile(path):
|
|
# 获取文件的 MIME 类型
|
|
mime_type, _ = mimetypes.guess_type(path)
|
|
if mime_type:
|
|
# 根据 MIME 类型判断是否为图像或视频
|
|
if mime_type.startswith('image'):
|
|
return 2
|
|
elif mime_type.startswith('video'):
|
|
return 3
|
|
return "文件"
|
|
|
|
return "未知类型"
|
|
|
|
class Fine_grained_T:
|
|
def __init__(self, onnx_model):
|
|
"""
|
|
Initialization.
|
|
|
|
Args:
|
|
onnx_model (str): Path to the ONNX model.
|
|
"""
|
|
if 'CUDAExecutionProvider' in ort.get_available_providers():
|
|
print("CUDA is available. ONNXRuntime will use GPU.")
|
|
self.session = ort.InferenceSession(
|
|
onnx_model,
|
|
providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
if ort.get_device() == "GPU"
|
|
else ["CPUExecutionProvider"],
|
|
)
|
|
else:
|
|
print("CUDA is not available. ONNXRuntime will use CPU.")
|
|
print(ort.get_device())
|
|
print("Fine is ready")
|
|
|
|
|
|
|
|
def inference(self, img, imagelist, labellist, scorelist, fine_thresh1,infsize):
|
|
with torch.no_grad():
|
|
for i, box in enumerate(imagelist):
|
|
if box[0] < 0:
|
|
box[0] = 0
|
|
if box[1] < 0:
|
|
box[1] = 0
|
|
if box[2] < 0:
|
|
box[2] = 0
|
|
if box[3] < 0:
|
|
box[3] = 0
|
|
X_left = int(box[1])-10
|
|
if X_left < 0:
|
|
X_left = 0
|
|
X_right = int(box[3])+10
|
|
if X_right>2159:
|
|
X_right = 2159
|
|
|
|
y_up = int(box[0])-10
|
|
if y_up <0 :
|
|
y_up = 0
|
|
|
|
y_down = int(box[2])+10
|
|
if y_down>3839:
|
|
y_down =3839
|
|
image = img.copy()[X_left:X_right, y_up:y_down, :]
|
|
image = Image.fromarray(image)
|
|
#image.save("/data5/Hatcher/data/count/testresult/11111.jpg")
|
|
#if gp.DebugPrint_flag:
|
|
#print(box)
|
|
#print(Img)
|
|
#if (int(box[3])-int(box[1])>gp.fine_thresh) and (int(box[2])-int(box[0])>gp.fine_thresh):
|
|
if ((int(box[3]) - int(box[1])) * (int(box[2]) - int(box[0]))) > fine_thresh1 :
|
|
image = image.resize((infsize, infsize)) # 假设模型需要 224x224 的输入
|
|
image = np.array(image).astype('float32')
|
|
image = (np.transpose(image, (2, 0, 1))) # 改变通道顺序
|
|
image = np.expand_dims(image, axis=0)
|
|
inputs = {self.session.get_inputs()[0].name: image}
|
|
outputs = self.session.run(None, inputs)
|
|
fineclass = np.argmax(outputs[0])
|
|
labellist[i] = fineclass
|
|
else:
|
|
pass
|
|
|
|
return imagelist, labellist, scorelist
|
|
def inferenceone(self, imagepath):
|
|
image = Image.open(imagepath)
|
|
image = image.resize((224, 224)) # 假设模型需要 224x224 的输入
|
|
image = np.array(image).astype('float32')
|
|
image = (np.transpose(image, (2, 0, 1))) # 改变通道顺序
|
|
image = np.expand_dims(image, axis=0)
|
|
inputs = {self.session.get_inputs()[0].name: image}
|
|
outputs = self.session.run(None, inputs)
|
|
fineclass = np.argmax(outputs[0])
|
|
#print(fineclass)
|
|
return(fineclass)
|
|
#print(fineclass)
|
|
class Detection:
|
|
def __init__(self,device,DetWeight,ChineseName):
|
|
self.device = device
|
|
self.img_size = 2176
|
|
self.conf_thres = 0.6
|
|
self.iou_thres = 0.4
|
|
self.stride = 1
|
|
self.weights = DetWeight
|
|
self.fine_thresh = 0.6
|
|
self.fineinfsize = 224
|
|
model = attempt_load(self.weights, device=self.device)
|
|
model.to(self.device).eval()
|
|
|
|
model.half()
|
|
self.model = model
|
|
self.names = model.module.names if hasattr(
|
|
model, 'module') else model.names
|
|
self.ChineseName = ChineseName
|
|
print("Detection is in ",self.device)
|
|
print("Detection is ready")
|
|
# while(1):
|
|
# pass
|
|
|
|
def preprocess(self, img):
|
|
|
|
img0 = img.copy()
|
|
#img = cv2.resize(img, (2000, 1000))
|
|
img = letterbox(img, new_shape=self.img_size, stride=self.stride)[0]
|
|
img = img[:, :, ::-1].transpose(2, 0, 1)
|
|
img = np.ascontiguousarray(img)
|
|
img = torch.from_numpy(img).to(self.device)
|
|
img = img.half() # ¡ã????¨¨
|
|
img /= 255.0 # ¨ª???1¨¦¨°??¡¥
|
|
if img.ndimension() == 3:
|
|
img = img.unsqueeze(0)
|
|
|
|
return img0, img
|
|
def AIlabelFine(self, imagename, fine_grained,saveroot,islabel):#imagename 是文件的完整名字
|
|
with torch.no_grad():
|
|
# t2 = time.time()
|
|
|
|
try:
|
|
t1 = time.time()
|
|
image = cv2.imread(imagename)
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("加载图片的时间是 ",(t2-t1)*1000)
|
|
except:
|
|
#os.remove(imagename)
|
|
print(imagename,"is broken")
|
|
return
|
|
if not type(image):
|
|
print("skip")
|
|
return
|
|
image_or = image
|
|
|
|
'''可视化注释
|
|
img_show = image_or.copy()
|
|
if (isinstance(img_show, np.ndarray)):
|
|
img_show = Image.fromarray(cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB))
|
|
draw = ImageDraw.Draw(img_show)
|
|
'''
|
|
try:
|
|
t1 = time.time()
|
|
img_detect = cv2.resize(image, (3840, 2160))
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("resize图片的时间是 ",(t2-t1)*1000)
|
|
except:
|
|
print(imagename+" is broken")
|
|
return
|
|
# t1 = time.time()
|
|
t1 = time.time()
|
|
image_scale, image = self.preprocess(img_detect)
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("预处理的时间是 ",(t2-t1)*1000)
|
|
# t2 = time.time()
|
|
# if DebugPrint_flag:
|
|
# tqdm.write("Detection data deal is ",(t2-t1)*1000)
|
|
t1 = time.time()
|
|
pred = self.model(image, augment=False)[0]
|
|
pred = pred.float()
|
|
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
|
|
# t3 = time.time()
|
|
# if gp.DebugPrint_flag:
|
|
# tqdm.write("Detection inference is ",(t3-t2)*1000)
|
|
boxes = []
|
|
scores = []
|
|
labels = []
|
|
for det in pred:
|
|
if det is not None and len(det):
|
|
det[:, :4] = scale_coords(
|
|
image.shape[2:], det[:, :4], image_scale.shape).round()
|
|
boxes.append(det[:, :4])
|
|
scores.append(det[:, 4])
|
|
labels.append((det[:, -1]))
|
|
boxes = boxes[0]
|
|
scores = scores[0]
|
|
labels = labels[0]
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("一次目标检测的时间是 ",(t2-t1)*1000)
|
|
# t5 = time.time()
|
|
# if gp.DebugPrint_flag:
|
|
# tqdm.write("Detection get result is ", (t5 - t3) * 1000)
|
|
t1 = time.time()
|
|
if len(labels) > 0:
|
|
ratio_x = (image_or.shape[1] / 3840.0)
|
|
ratio_y = (image_or.shape[0] / 2160.0)
|
|
for i in range(len(boxes)):
|
|
boxes[i][0] = boxes[i][0] * ratio_x
|
|
boxes[i][2] = boxes[i][2] * ratio_x
|
|
boxes[i][1] = boxes[i][1] * ratio_y
|
|
boxes[i][3] = boxes[i][3] * ratio_y
|
|
if boxes[i][0] < 0:
|
|
boxes[i][0] = 0
|
|
if boxes[i][1] < 0:
|
|
boxes[i][1] = 0
|
|
if boxes[i][2] < 0:
|
|
boxes[i][2] = 0
|
|
if boxes[i][3] < 0:
|
|
boxes[i][3] = 0
|
|
for i in range(len(labels)):
|
|
labels[i] = 0
|
|
boxes_fine, labels_fine, scores_fine = fine_grained.inference(image_or, boxes, labels, scores, self.fine_thresh,self.fineinfsize)
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("细粒度识别的时间是 ",(t2-t1)*1000)
|
|
# t4 = time.time()
|
|
# if gp.DebugPrint_flag:
|
|
# tqdm.write("Detection final is ", (t4 - t3) * 1000)
|
|
'''不进行可视化
|
|
for i in range(len(scores)):
|
|
xyxy = boxes_fine[i]
|
|
conf = scores_fine[i]
|
|
cls = labels_fine[i]
|
|
print(int(cls))
|
|
label = str(self.ChineseName[int(cls)])
|
|
#label = f'{int(cls)} {conf:.2f}'
|
|
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
|
|
colors = tuple(int((p * (cls ** 2 - cls + 1)) % 255) for p in palette)
|
|
draw = plot_one_box(xyxy, draw, image_or.shape[1], image_or.shape[0], label=label, color=colors,
|
|
line_thickness=int(3 * gp.radioShowtext))
|
|
savepath = os.path.join(saveroot,os.path.basename(imagename))
|
|
img_show = cv2.cvtColor(np.asarray(img_show), cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite(savepath,img_show)
|
|
'''
|
|
if islabel:
|
|
# for i in range(len(scores)):
|
|
|
|
# cls = labels_fine[i]
|
|
# print(int(cls))
|
|
# labels_fine[i] = str(self.ChineseName[int(cls)])
|
|
#print()
|
|
t1 = time.time()
|
|
savepath = os.path.join(saveroot,os.path.basename(imagename))
|
|
xml_outputpath = savepath.replace(".jpg", ".xml").replace(".JPG", ".xml")
|
|
make_xml(imagename, image_or.shape[0], image_or.shape[1], boxes_fine, labels_fine, scores_fine, xml_outputpath)
|
|
t2 = time.time()
|
|
if DebugPrint_flag:
|
|
print("结果存储的时间是 ",(t2-t1)*1000)
|
|
|
|
else:
|
|
pass
|
|
'''不保存数据
|
|
savepath = os.path.join(saveroot,os.path.basename(imagename))
|
|
shutil.copy(imagename, savepath)
|
|
'''
|
|
def get_all_image_paths(folder_path):
|
|
|
|
# Define common image file extensions
|
|
image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.tif')
|
|
|
|
# Initialize an empty list to store the image paths
|
|
image_paths = []
|
|
|
|
# Walk through the folder and its subfolders
|
|
for root, dirs, files in tqdm(os.walk(folder_path)):
|
|
for file in files:
|
|
if file.lower().endswith(image_extensions):
|
|
# Get the absolute path of the image file
|
|
image_path = os.path.join(root, file)
|
|
image_paths.append(image_path)
|
|
|
|
return image_paths
|
|
def process_image(filepath, Fine, saveroot, islabel, detection):
|
|
detection.AIlabelFine(filepath, Fine, saveroot, islabel)
|
|
if __name__ == '__main__':
|
|
# Fine = Fine_grained_T("/data5/Train_Models/mmcls_onnx/CMB_YNDC_20240115_29/convnext-base_4xb32_cmb200/end2end.onnx")
|
|
# result = Fine.inferenceone("/data5/CAS_Project/CAS_ProjectBird/CMB/CMB_YNDC/CMB_YNDC_20231129_30/OrginalData/Bird/Bird___BJCH014_000447.jpg")
|
|
"""
|
|
测试
|
|
-dw /data5/Train_Models/Detection_weigth/3090/202306/20230619-131704-pixel5000-BJNHZ_Person/weights/best.pt
|
|
-cw /data5/Train_Models/mmcls_onnx/CMB_BJNHZ_20230910_144/convnext-base_4xb32_cmb200/end2end.onnx
|
|
-c /data5/CAS_Project/CAS_ProjectBird/CMB/CMB_BJNHZ/CMB_BJNHZ_20230910_144/birdListbirdCNlist.txt
|
|
-m "ai"
|
|
-s "/data5/Hatcher/data/detection/test1"
|
|
-o "/data5/Hatcher/data/detection/testresult"
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
Det_weights,Fine_weights,Fine_class,infmode,sourcepath,saveroot = parse_args()
|
|
if not os.path.exists(saveroot):
|
|
os.makedirs(saveroot)
|
|
# with open(Fine_class, 'r', encoding='utf-8') as file:
|
|
# bird_list = file.readlines()
|
|
with open(Fine_class, 'r', encoding='utf-8') as file:
|
|
content = file.read().strip() # 读取文件内容并去除首尾空格
|
|
bird_list = content.split(';') # 使用分号分割字符串
|
|
|
|
#bird_list = [bird.strip() for bird in bird_list]
|
|
print(bird_list)
|
|
print(saveroot)
|
|
devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
|
detection = Detection(devices,Det_weights,bird_list)
|
|
Fine = Fine_grained_T(Fine_weights)
|
|
if not os.path.exists(saveroot):
|
|
os.makedirs(saveroot)
|
|
fileType = identify_path_type(sourcepath)
|
|
if infmode == "ai":
|
|
islabel = 1
|
|
if (fileType == 2):
|
|
print("开始处理单张图像数据")
|
|
#缺一个处理单个的函数
|
|
detection.AIlabelFine(sourcepath,Fine,saveroot,islabel)
|
|
if (fileType == 1): #说明是文件夹
|
|
print("开始处理文件夹批量数据")
|
|
fileList = os.listdir(sourcepath)
|
|
for filepath in tqdm(fileList):
|
|
sourcePath = os.path.join(sourcepath,filepath)
|
|
#缺一个处理单个的函数
|
|
#print(sourcePath)
|
|
detection.AIlabelFine(sourcePath,Fine,saveroot,islabel)
|
|
if (fileType == 4): #说明是文件夹
|
|
print("开始多线程处理复杂文件夹批量数据")
|
|
fileList = get_all_image_paths(sourcepath)
|
|
#print(fileList)
|
|
for filepath in tqdm(fileList):
|
|
#sourcePath = os.path.join(sourcepath,filepath)
|
|
#缺一个处理单个的函数
|
|
#print(sourcePath)
|
|
with ThreadPoolExecutor(max_workers=12) as executor:
|
|
futures = {executor.submit(process_image, filepath, Fine, saveroot, islabel, detection): filepath for filepath in fileList}
|
|
|
|
for future in tqdm(as_completed(futures), total=len(futures)):
|
|
try:
|
|
future.result()
|
|
except Exception as exc:
|
|
print(f'File {futures[future]} generated an exception: {exc}')
|
|
# t1 = time.time()
|
|
# detection.AIlabelFine(filepath,Fine,saveroot,islabel)
|
|
# t2 = time.time()
|
|
# if DebugPrint_flag:
|
|
# print("推理一张图片的时间是 ",(t2-t1)*1000)
|
|
if (fileType == "always"):
|
|
print("开始循环处理输入数据")
|
|
while(1):
|
|
sourcepath = input("请输入测试文件路径： ")
|
|
if sourcepath == "q":
|
|
break
|
|
detection.AIlabelFine(sourcepath,Fine,saveroot,islabel)
|
|
#缺一个处理单个的函数 |