""" 测试代码： 1.-dw 输入目标检测权重，为pt文件，必要选项 2.-cw 输入细粒度识别权重，为onnx文件，必要选项 3.-c 输入细粒度识别名录，为txt文件，中文名字，每个中文用引号表示，中间用“，”隔开 4.-m 系统执行模式，分两种，一种为推理图片，直接生成结果“inf”，一种为AI label，生成xml文件，“ai” 5.-s 源文件路径，可以是文件夹也可以是图片文件，如果是AI label，则在同级目录下面生成一个xml文件。，可以一直输入单张图片地址进行推理“always”，则源文件路径参数失效，需要输入1即可，也可以是超大分辨率的处理，返回的参数是7，要求文件夹目录必须有long字眼 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 0: 路径不存在 1: 普通文件夹（不包含子文件夹） 2: 图像文件 3: 视频文件 4: 包含子文件夹的文件夹 6: TXT文本文件 7: 名称包含"long"的文件夹 "always": 特殊标识符 "文件": 其他类型文件 "未知类型": 无法识别的类型 """ 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 try: import thop # for FLOPs computation except ImportError: thop = None import numpy as np import pkg_resources as pkg 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") 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): 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])) 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) 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: 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 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): # 检查路径是否存在 print("原始路径为：{}".format(path)) if path =="always": return "always" if not os.path.exists(path): return 0 if 'long' in path.lower(): return 7 # 检查是否为文件夹 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 : # 定义归一化参数 img_norm_cfg = dict( mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True ) image = image.resize((infsize, infsize)) image = np.array(image).astype('float32') if img_norm_cfg['to_rgb']: image = image[..., ::-1] # BGR转RGB image = np.transpose(image, (2, 0, 1)) # HWC转CHW mean = np.array(img_norm_cfg['mean']).reshape(-1, 1, 1) std = np.array(img_norm_cfg['std']).reshape(-1, 1, 1) image = (image - mean) / std image = np.expand_dims(image, axis=0) image = image.astype('float32') # 模型推理 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.3#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 process_long_image(self, img_path, fine_grained, save_root): image = cv2.imread(img_path) if image is None: print(f"Failed to load image: {img_path}") return original_height, original_width = image.shape[:2] # 计算resize后的尺寸 target_width = 3840 num_splits = max(1, round(original_width / target_width)) new_width = target_width * num_splits scale_factor = new_width / original_width new_height = int(original_height * scale_factor) # resize图像 resized_image = cv2.resize(image, (new_width, new_height)) all_boxes = [] all_labels = [] all_scores = [] for i in range(num_splits): start_x = i * target_width end_x = (i + 1) * target_width sub_image = resized_image[:, start_x:end_x] # 检测处理 img_detect = cv2.resize(sub_image, (3840, 2160)) detect_height, detect_width = img_detect.shape[:2] image_scale, image = self.preprocess(img_detect) with torch.no_grad(): pred = self.model(image, augment=False)[0] pred = pred.float() pred = non_max_suppression(pred, self.conf_thres, self.iou_thres) for det in pred: if det is not None and len(det): # 先将坐标从模型输出转换到3840x2160 det[:, :4] = scale_coords( image.shape[2:], det[:, :4], img_detect.shape).round() # 计算从3840x2160到子图实际大小的缩放比例 width_scale = sub_image.shape[1] / detect_width height_scale = sub_image.shape[0] / detect_height # 转换坐标到子图实际大小 boxes = det[:, :4].cpu().numpy() boxes[:, [0, 2]] *= width_scale boxes[:, [1, 3]] *= height_scale # 添加偏移量 boxes[:, [0, 2]] += start_x scores = det[:, 4].cpu().numpy() labels = det[:, -1].cpu().numpy() if len(labels) > 0: boxes_fine, labels_fine, scores_fine = fine_grained.inference( resized_image, boxes, labels, scores, self.fine_thresh, self.fineinfsize) all_boxes.extend(boxes_fine) all_labels.extend(labels_fine) all_scores.extend(scores_fine) # 创建可视化结果 result_image = resized_image.copy() # 添加调试信息 for i, box in enumerate(all_boxes): print(f"Box {i}: {box}") for box, label, score in zip(all_boxes, all_labels, all_scores): box = [int(b) for b in box] color = (0, 255, 0) # 确保坐标不超出图像边界 x1 = max(0, min(box[0], result_image.shape[1]-1)) y1 = max(0, min(box[1], result_image.shape[0]-1)) x2 = max(0, min(box[2], result_image.shape[1]-1)) y2 = max(0, min(box[3], result_image.shape[0]-1)) cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2) label_text = "Oriental Stork"#f"{self.ChineseName[int(label)]} {score:.2f}" cv2.putText(result_image, label_text, (x1, max(y1-10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) if not os.path.exists(save_root): os.makedirs(save_root) output_path = os.path.join(save_root, os.path.basename(img_path)) cv2.imwrite(output_path, result_image) xml_path = output_path.replace('.jpg', '.xml').replace('.JPG', '.xml') make_xml(img_path, resized_image.shape[0], resized_image.shape[1], all_boxes, all_labels, all_scores, xml_path) 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 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() # 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) print("处理模式为{}".format(str(fileType))) 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) t1 = time.time() detection.AIlabelFine(filepath,Fine,saveroot,islabel) t2 = time.time() if DebugPrint_flag: print("推理一张图片的时间是 ",(t2-t1)*1000) if (fileType == 6): #说明是文件路径的txt print("开始处理txt文件中的批量数据") with open(sourcepath, 'r') as file: image_paths = [line.strip() for line in file.readlines()] fileList = image_paths #print(fileList) for filepath in tqdm(fileList): #sourcePath = os.path.join(sourcepath,filepath) #缺一个处理单个的函数 #print(sourcePath) t1 = time.time() detection.AIlabelFine(filepath,Fine,saveroot,islabel) t2 = time.time() if DebugPrint_flag: print("推理一张图片的时间是 ",(t2-t1)*1000) if (fileType == 7): #说明是文件路径的txt print("开始处理长图像数据") parent_dir = os.path.dirname(sourcepath) save_root = os.path.join(parent_dir, 'output_') if not os.path.exists(save_root): os.makedirs(save_root) # 处理文件夹中的所有图像 for img_name in tqdm(os.listdir(sourcepath)): if img_name.lower().endswith(('.jpg', '.jpeg', '.png')): img_path = os.path.join(sourcepath, img_name) detection.process_long_image(img_path, Fine, save_root) if (fileType == "always"): print("开始循环处理输入数据") while(1): sourcepath = input("请输入测试文件路径： ") if sourcepath == "q": break detection.AIlabelFine(sourcepath,Fine,saveroot,islabel) #缺一个处理单个的函数