Xlittlestar/faster_rcnn_coins_detection
0
1import math2from typing import List, Tuple, Dict, Optional3 4import torch5from torch import nn, Tensor6import torchvision7 8from .image_list import ImageList9 10 11@torch.jit.unused12def _resize_image_onnx(image, self_min_size, self_max_size):13 # type: (Tensor, float, float) -> Tensor14 from torch.onnx import operators15 im_shape = operators.shape_as_tensor(image)[-2:]16 min_size = torch.min(im_shape).to(dtype=torch.float32)17 max_size = torch.max(im_shape).to(dtype=torch.float32)18 scale_factor = torch.min(self_min_size / min_size, self_max_size / max_size)19 20 image = torch.nn.functional.interpolate(21 image[None], scale_factor=scale_factor, mode="bilinear", recompute_scale_factor=True,22 align_corners=False)[0]23 24 return image25 26 27def _resize_image(image, self_min_size, self_max_size):28 # type: (Tensor, float, float) -> Tensor29 im_shape = torch.tensor(image.shape[-2:])30 min_size = float(torch.min(im_shape)) # 获取高宽中的最小值31 max_size = float(torch.max(im_shape)) # 获取高宽中的最大值32 scale_factor = self_min_size / min_size # 根据指定最小边长和图片最小边长计算缩放比例33 34 # 如果使用该缩放比例计算的图片最大边长大于指定的最大边长35 if max_size * scale_factor > self_max_size:36 scale_factor = self_max_size / max_size # 将缩放比例设为指定最大边长和图片最大边长之比37 38 # interpolate利用插值的方法缩放图片39 # image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]40 # bilinear只支持4D Tensor41 image = torch.nn.functional.interpolate(42 image[None], scale_factor=scale_factor, mode="bilinear", recompute_scale_factor=True,43 align_corners=False)[0]44 45 return image46 47 48class GeneralizedRCNNTransform(nn.Module):49 """50 Performs input / target transformation before feeding the data to a GeneralizedRCNN51 model.52 53 The transformations it perform are:54 - input normalization (mean subtraction and std division)55 - input / target resizing to match min_size / max_size56 57 It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets58 """59 60 def __init__(self, min_size, max_size, image_mean, image_std):61 super(GeneralizedRCNNTransform, self).__init__()62 if not isinstance(min_size, (list, tuple)):63 min_size = (min_size,)64 self.min_size = min_size # 指定图像的最小边长范围65 self.max_size = max_size # 指定图像的最大边长范围66 self.image_mean = image_mean # 指定图像在标准化处理中的均值67 self.image_std = image_std # 指定图像在标准化处理中的方差68 69 def normalize(self, image):70 """标准化处理"""71 dtype, device = image.dtype, image.device72 mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)73 std = torch.as_tensor(self.image_std, dtype=dtype, device=device)74 # [:, None, None]: shape [3] -> [3, 1, 1]75 return (image - mean[:, None, None]) / std[:, None, None]76 77 def torch_choice(self, k):78 # type: (List[int]) -> int79 """80 Implements `random.choice` via torch ops so it can be compiled with81 TorchScript. Remove if https://github.com/pytorch/pytorch/issues/2580382 is fixed.83 """84 index = int(torch.empty(1).uniform_(0., float(len(k))).item())85 return k[index]86 87 def resize(self, image, target):88 # type: (Tensor, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]89 """90 将图片缩放到指定的大小范围内，并对应缩放bboxes信息91 Args:92 image: 输入的图片93 target: 输入图片的相关信息（包括bboxes信息）94 95 Returns:96 image: 缩放后的图片97 target: 缩放bboxes后的图片相关信息98 """99 # image shape is [channel, height, width]100 h, w = image.shape[-2:]101 102 if self.training:103 size = float(self.torch_choice(self.min_size)) # 指定输入图片的最小边长,注意是self.min_size不是min_size104 else:105 # FIXME assume for now that testing uses the largest scale106 size = float(self.min_size[-1]) # 指定输入图片的最小边长,注意是self.min_size不是min_size107 108 if torchvision._is_tracing():109 image = _resize_image_onnx(image, size, float(self.max_size))110 else:111 image = _resize_image(image, size, float(self.max_size))112 113 if target is None:114 return image, target115 116 bbox = target["boxes"]117 # 根据图像的缩放比例来缩放bbox118 bbox = resize_boxes(bbox, [h, w], image.shape[-2:])119 target["boxes"] = bbox120 121 return image, target122 123 # _onnx_batch_images() is an implementation of124 # batch_images() that is supported by ONNX tracing.125 @torch.jit.unused126 def _onnx_batch_images(self, images, size_divisible=32):127 # type: (List[Tensor], int) -> Tensor128 max_size = []129 for i in range(images[0].dim()):130 max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)131 max_size.append(max_size_i)132 stride = size_divisible133 max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)134 max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)135 max_size = tuple(max_size)136 137 # work around for138 # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)139 # which is not yet supported in onnx140 padded_imgs = []141 for img in images:142 padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]143 padded_img = torch.nn.functional.pad(img, [0, padding[2], 0, padding[1], 0, padding[0]])144 padded_imgs.append(padded_img)145 146 return torch.stack(padded_imgs)147 148 def max_by_axis(self, the_list):149 # type: (List[List[int]]) -> List[int]150 maxes = the_list[0]151 for sublist in the_list[1:]:152 for index, item in enumerate(sublist):153 maxes[index] = max(maxes[index], item)154 return maxes155 156 def batch_images(self, images, size_divisible=32):157 # type: (List[Tensor], int) -> Tensor158 """159 将一批图像打包成一个batch返回（注意batch中每个tensor的shape是相同的）160 Args:161 images: 输入的一批图片162 size_divisible: 将图像高和宽调整到该数的整数倍163 164 Returns:165 batched_imgs: 打包成一个batch后的tensor数据166 """167 168 if torchvision._is_tracing():169 # batch_images() does not export well to ONNX170 # call _onnx_batch_images() instead171 return self._onnx_batch_images(images, size_divisible)172 173 # 分别计算一个batch中所有图片中的最大channel, height, width174 max_size = self.max_by_axis([list(img.shape) for img in images])175 176 stride = float(size_divisible)177 # max_size = list(max_size)178 # 将height向上调整到stride的整数倍179 max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)180 # 将width向上调整到stride的整数倍181 max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)182 183 # [batch, channel, height, width]184 batch_shape = [len(images)] + max_size185 186 # 创建shape为batch_shape且值全部为0的tensor187 batched_imgs = images[0].new_full(batch_shape, 0)188 for img, pad_img in zip(images, batched_imgs):189 # 将输入images中的每张图片复制到新的batched_imgs的每张图片中，对齐左上角，保证bboxes的坐标不变190 # 这样保证输入到网络中一个batch的每张图片的shape相同191 # copy_: Copies the elements from src into self tensor and returns self192 pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)193 194 return batched_imgs195 196 def postprocess(self,197 result, # type: List[Dict[str, Tensor]]198 image_shapes, # type: List[Tuple[int, int]]199 original_image_sizes # type: List[Tuple[int, int]]200 ):201 # type: (...) -> List[Dict[str, Tensor]]202 """203 对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）204 Args:205 result: list(dict), 网络的预测结果, len(result) == batch_size206 image_shapes: list(torch.Size), 图像预处理缩放后的尺寸, len(image_shapes) == batch_size207 original_image_sizes: list(torch.Size), 图像的原始尺寸, len(original_image_sizes) == batch_size208 209 Returns:210 211 """212 if self.training:213 return result214 215 # 遍历每张图片的预测信息，将boxes信息还原回原尺度216 for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):217 boxes = pred["boxes"]218 boxes = resize_boxes(boxes, im_s, o_im_s) # 将bboxes缩放回原图像尺度上219 result[i]["boxes"] = boxes220 return result221 222 def __repr__(self):223 """自定义输出实例化对象的信息，可通过print打印实例信息"""224 format_string = self.__class__.__name__ + '('225 _indent = '\n '226 format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std)227 format_string += "{0}Resize(min_size={1}, max_size={2}, mode='bilinear')".format(_indent, self.min_size,228 self.max_size)229 format_string += '\n)'230 return format_string231 232 def forward(self,233 images, # type: List[Tensor]234 targets=None # type: Optional[List[Dict[str, Tensor]]]235 ):236 # type: (...) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]237 images = [img for img in images]238 for i in range(len(images)):239 image = images[i]240 target_index = targets[i] if targets is not None else None241 242 if image.dim() != 3:243 raise ValueError("images is expected to be a list of 3d tensors "244 "of shape [C, H, W], got {}".format(image.shape))245 image = self.normalize(image) # 对图像进行标准化处理246 image, target_index = self.resize(image, target_index) # 对图像和对应的bboxes缩放到指定范围247 images[i] = image248 if targets is not None and target_index is not None:249 targets[i] = target_index250 251 # 记录resize后的图像尺寸252 image_sizes = [img.shape[-2:] for img in images]253 images = self.batch_images(images) # 将images打包成一个batch254 image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])255 256 for image_size in image_sizes:257 assert len(image_size) == 2258 image_sizes_list.append((image_size[0], image_size[1]))259 260 image_list = ImageList(images, image_sizes_list)261 return image_list, targets262 263 264def resize_boxes(boxes, original_size, new_size):265 # type: (Tensor, List[int], List[int]) -> Tensor266 """267 将boxes参数根据图像的缩放情况进行相应缩放268 269 Arguments:270 original_size: 图像缩放前的尺寸271 new_size: 图像缩放后的尺寸272 """273 ratios = [274 torch.tensor(s, dtype=torch.float32, device=boxes.device) /275 torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)276 for s, s_orig in zip(new_size, original_size)277 ]278 ratios_height, ratios_width = ratios279 # Removes a tensor dimension, boxes [minibatch, 4]280 # Returns a tuple of all slices along a given dimension, already without it.281 xmin, ymin, xmax, ymax = boxes.unbind(1)282 xmin = xmin * ratios_width283 xmax = xmax * ratios_width284 ymin = ymin * ratios_height285 ymax = ymax * ratios_height286 return torch.stack((xmin, ymin, xmax, ymax), dim=1)