line_predictor.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348
  1. from typing import Any, Optional
  2. import torch
  3. from torch import nn
  4. from torchvision.ops import MultiScaleRoIAlign
  5. from libs.vision_libs.ops import misc as misc_nn_ops
  6. from libs.vision_libs.transforms._presets import ObjectDetection
  7. from .roi_heads import RoIHeads
  8. from libs.vision_libs.models._api import register_model, Weights, WeightsEnum
  9. from libs.vision_libs.models._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES
  10. from libs.vision_libs.models._utils import _ovewrite_value_param, handle_legacy_interface
  11. from libs.vision_libs.models.resnet import resnet50, ResNet50_Weights
  12. from libs.vision_libs.models.detection._utils import overwrite_eps
  13. from libs.vision_libs.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers
  14. from libs.vision_libs.models.detection.faster_rcnn import FasterRCNN, TwoMLPHead, FastRCNNPredictor
  15. from models.config.config_tool import read_yaml
  16. import numpy as np
  17. import torch.nn.functional as F
  18. FEATURE_DIM = 8
  19. def non_maximum_suppression(a):
  20. ap = F.max_pool2d(a, 3, stride=1, padding=1)
  21. mask = (a == ap).float().clamp(min=0.0)
  22. return a * mask
  23. class Bottleneck1D(nn.Module):
  24. def __init__(self, inplanes, outplanes):
  25. super(Bottleneck1D, self).__init__()
  26. planes = outplanes // 2
  27. self.op = nn.Sequential(
  28. nn.BatchNorm1d(inplanes),
  29. nn.ReLU(inplace=True),
  30. nn.Conv1d(inplanes, planes, kernel_size=1),
  31. nn.BatchNorm1d(planes),
  32. nn.ReLU(inplace=True),
  33. nn.Conv1d(planes, planes, kernel_size=3, padding=1),
  34. nn.BatchNorm1d(planes),
  35. nn.ReLU(inplace=True),
  36. nn.Conv1d(planes, outplanes, kernel_size=1),
  37. )
  38. def forward(self, x):
  39. return x + self.op(x)
  40. class LineRCNNPredictor(nn.Module):
  41. def __init__(self, cfg):
  42. super().__init__()
  43. # self.backbone = backbone
  44. # self.cfg = read_yaml(cfg)
  45. # self.cfg = read_yaml(r'./config/wireframe.yaml')
  46. self.cfg = cfg
  47. self.n_pts0 = self.cfg['n_pts0']
  48. self.n_pts1 = self.cfg['n_pts1']
  49. self.n_stc_posl = self.cfg['n_stc_posl']
  50. self.dim_loi = self.cfg['dim_loi']
  51. self.use_conv = self.cfg['use_conv']
  52. self.dim_fc = self.cfg['dim_fc']
  53. self.n_out_line = self.cfg['n_out_line']
  54. self.n_out_junc = self.cfg['n_out_junc']
  55. self.loss_weight = self.cfg['loss_weight']
  56. self.n_dyn_junc = self.cfg['n_dyn_junc']
  57. self.eval_junc_thres = self.cfg['eval_junc_thres']
  58. self.n_dyn_posl = self.cfg['n_dyn_posl']
  59. self.n_dyn_negl = self.cfg['n_dyn_negl']
  60. self.n_dyn_othr = self.cfg['n_dyn_othr']
  61. self.use_cood = self.cfg['use_cood']
  62. self.use_slop = self.cfg['use_slop']
  63. self.n_stc_negl = self.cfg['n_stc_negl']
  64. self.head_size = self.cfg['head_size']
  65. self.num_class = sum(sum(self.head_size, []))
  66. self.head_off = np.cumsum([sum(h) for h in self.head_size])
  67. lambda_ = torch.linspace(0, 1, self.n_pts0)[:, None]
  68. self.register_buffer("lambda_", lambda_)
  69. self.do_static_sampling = self.n_stc_posl + self.n_stc_negl > 0
  70. self.fc1 = nn.Conv2d(256, self.dim_loi, 1)
  71. scale_factor = self.n_pts0 // self.n_pts1
  72. if self.use_conv:
  73. self.pooling = nn.Sequential(
  74. nn.MaxPool1d(scale_factor, scale_factor),
  75. Bottleneck1D(self.dim_loi, self.dim_loi),
  76. )
  77. self.fc2 = nn.Sequential(
  78. nn.ReLU(inplace=True), nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, 1)
  79. )
  80. else:
  81. self.pooling = nn.MaxPool1d(scale_factor, scale_factor)
  82. self.fc2 = nn.Sequential(
  83. nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, self.dim_fc),
  84. nn.ReLU(inplace=True),
  85. nn.Linear(self.dim_fc, self.dim_fc),
  86. nn.ReLU(inplace=True),
  87. nn.Linear(self.dim_fc, 1),
  88. )
  89. self.loss = nn.BCEWithLogitsLoss(reduction="none")
  90. def forward(self, inputs, features, targets=None):
  91. # outputs, features = input
  92. # for out in outputs:
  93. # print(f'out:{out.shape}')
  94. # outputs=merge_features(outputs,100)
  95. batch, channel, row, col = inputs.shape
  96. # print(f'outputs:{inputs.shape}')
  97. # print(f'batch:{batch}, channel:{channel}, row:{row}, col:{col}')
  98. if targets is not None:
  99. self.training = True
  100. # print(f'target:{targets}')
  101. wires_targets = [t["wires"] for t in targets]
  102. # print(f'wires_target:{wires_targets}')
  103. # 提取所有 'junc_map', 'junc_offset', 'line_map' 的张量
  104. junc_maps = [d["junc_map"] for d in wires_targets]
  105. junc_offsets = [d["junc_offset"] for d in wires_targets]
  106. line_maps = [d["line_map"] for d in wires_targets]
  107. junc_map_tensor = torch.stack(junc_maps, dim=0)
  108. junc_offset_tensor = torch.stack(junc_offsets, dim=0)
  109. line_map_tensor = torch.stack(line_maps, dim=0)
  110. wires_meta = {
  111. "junc_map": junc_map_tensor,
  112. "junc_offset": junc_offset_tensor,
  113. # "line_map": line_map_tensor,
  114. }
  115. else:
  116. self.training = False
  117. t = {
  118. "junc_coords": torch.zeros(1, 2),
  119. "jtyp": torch.zeros(1, dtype=torch.uint8),
  120. "line_pos_idx": torch.zeros(2, 2, dtype=torch.uint8),
  121. "line_neg_idx": torch.zeros(2, 2, dtype=torch.uint8),
  122. "junc_map": torch.zeros([1, 1, 128, 128]),
  123. "junc_offset": torch.zeros([1, 1, 2, 128, 128]),
  124. }
  125. wires_targets = [t for b in range(inputs.size(0))]
  126. wires_meta = {
  127. "junc_map": torch.zeros([1, 1, 128, 128]),
  128. "junc_offset": torch.zeros([1, 1, 2, 128, 128]),
  129. }
  130. T = wires_meta.copy()
  131. n_jtyp = T["junc_map"].shape[1]
  132. offset = self.head_off
  133. result = {}
  134. for stack, output in enumerate([inputs]):
  135. output = output.transpose(0, 1).reshape([-1, batch, row, col]).contiguous()
  136. # print(f"Stack {stack} output shape: {output.shape}") # 打印每层的输出形状
  137. jmap = output[0: offset[0]].reshape(n_jtyp, 2, batch, row, col)
  138. lmap = output[offset[0]: offset[1]].squeeze(0)
  139. joff = output[offset[1]: offset[2]].reshape(n_jtyp, 2, batch, row, col)
  140. if stack == 0:
  141. result["preds"] = {
  142. "jmap": jmap.permute(2, 0, 1, 3, 4).softmax(2)[:, :, 1],
  143. "lmap": lmap.sigmoid(),
  144. "joff": joff.permute(2, 0, 1, 3, 4).sigmoid() - 0.5,
  145. }
  146. # visualize_feature_map(jmap[0, 0], title=f"jmap - Stack {stack}")
  147. # visualize_feature_map(lmap, title=f"lmap - Stack {stack}")
  148. # visualize_feature_map(joff[0, 0], title=f"joff - Stack {stack}")
  149. h = result["preds"]
  150. # print(f'features shape:{features.shape}')
  151. x = self.fc1(features)
  152. # print(f'x:{x.shape}')
  153. n_batch, n_channel, row, col = x.shape
  154. # print(f'n_batch:{n_batch}, n_channel:{n_channel}, row:{row}, col:{col}')
  155. xs, ys, fs, ps, idx, jcs = [], [], [], [], [0], []
  156. for i, meta in enumerate(wires_targets):
  157. p, label, feat, jc = self.sample_lines(
  158. meta, h["jmap"][i], h["joff"][i],
  159. )
  160. # print(f"p.shape:{p.shape},label:{label.shape},feat:{feat.shape},jc:{len(jc)}")
  161. ys.append(label)
  162. if self.training and self.do_static_sampling:
  163. p = torch.cat([p, meta["lpre"]])
  164. feat = torch.cat([feat, meta["lpre_feat"]])
  165. ys.append(meta["lpre_label"])
  166. del jc
  167. else:
  168. jcs.append(jc)
  169. ps.append(p)
  170. fs.append(feat)
  171. p = p[:, 0:1, :] * self.lambda_ + p[:, 1:2, :] * (1 - self.lambda_) - 0.5
  172. p = p.reshape(-1, 2) # [N_LINE x N_POINT, 2_XY]
  173. px, py = p[:, 0].contiguous(), p[:, 1].contiguous()
  174. px0 = px.floor().clamp(min=0, max=127)
  175. py0 = py.floor().clamp(min=0, max=127)
  176. px1 = (px0 + 1).clamp(min=0, max=127)
  177. py1 = (py0 + 1).clamp(min=0, max=127)
  178. px0l, py0l, px1l, py1l = px0.long(), py0.long(), px1.long(), py1.long()
  179. # xp: [N_LINE, N_CHANNEL, N_POINT]
  180. xp = (
  181. (
  182. x[i, :, px0l, py0l] * (px1 - px) * (py1 - py)
  183. + x[i, :, px1l, py0l] * (px - px0) * (py1 - py)
  184. + x[i, :, px0l, py1l] * (px1 - px) * (py - py0)
  185. + x[i, :, px1l, py1l] * (px - px0) * (py - py0)
  186. )
  187. .reshape(n_channel, -1, self.n_pts0)
  188. .permute(1, 0, 2)
  189. )
  190. xp = self.pooling(xp)
  191. # print(f'xp.shape:{xp.shape}')
  192. xs.append(xp)
  193. idx.append(idx[-1] + xp.shape[0])
  194. # print(f'idx__:{idx}')
  195. x, y = torch.cat(xs), torch.cat(ys)
  196. f = torch.cat(fs)
  197. x = x.reshape(-1, self.n_pts1 * self.dim_loi)
  198. # print("Weight dtype:", self.fc2.weight.dtype)
  199. x = torch.cat([x, f], 1)
  200. # print("Input dtype:", x.dtype)
  201. x = x.to(dtype=torch.float32)
  202. # print("Input dtype1:", x.dtype)
  203. x = self.fc2(x).flatten()
  204. # return x,idx,jcs,n_batch,ps,self.n_out_line,self.n_out_junc
  205. return x, y, idx, jcs, n_batch, ps, self.n_out_line, self.n_out_junc
  206. # if mode != "training":
  207. # self.inference(x, idx, jcs, n_batch, ps)
  208. # return result
  209. def sample_lines(self, meta, jmap, joff):
  210. device = jmap.device
  211. with torch.no_grad():
  212. junc = meta["junc_coords"].to(device) # [N, 2]
  213. jtyp = meta["jtyp"].to(device) # [N]
  214. Lpos = meta["line_pos_idx"].to(device)
  215. Lneg = meta["line_neg_idx"].to(device)
  216. n_type = jmap.shape[0]
  217. jmap = non_maximum_suppression(jmap).reshape(n_type, -1)
  218. joff = joff.reshape(n_type, 2, -1)
  219. max_K = self.n_dyn_junc // n_type
  220. N = len(junc)
  221. # if mode != "training":
  222. if not self.training:
  223. K = min(int((jmap > self.eval_junc_thres).float().sum().item()), max_K)
  224. else:
  225. K = min(int(N * 2 + 2), max_K)
  226. if K < 2:
  227. K = 2
  228. device = jmap.device
  229. # index: [N_TYPE, K]
  230. score, index = torch.topk(jmap, k=K)
  231. y = (index // 128).float() + torch.gather(joff[:, 0], 1, index) + 0.5
  232. x = (index % 128).float() + torch.gather(joff[:, 1], 1, index) + 0.5
  233. # xy: [N_TYPE, K, 2]
  234. xy = torch.cat([y[..., None], x[..., None]], dim=-1)
  235. xy_ = xy[..., None, :]
  236. del x, y, index
  237. # dist: [N_TYPE, K, N]
  238. dist = torch.sum((xy_ - junc) ** 2, -1)
  239. cost, match = torch.min(dist, -1)
  240. # xy: [N_TYPE * K, 2]
  241. # match: [N_TYPE, K]
  242. for t in range(n_type):
  243. match[t, jtyp[match[t]] != t] = N
  244. match[cost > 1.5 * 1.5] = N
  245. match = match.flatten()
  246. _ = torch.arange(n_type * K, device=device)
  247. u, v = torch.meshgrid(_, _)
  248. u, v = u.flatten(), v.flatten()
  249. up, vp = match[u], match[v]
  250. label = Lpos[up, vp]
  251. # if mode == "training":
  252. if self.training:
  253. c = torch.zeros_like(label, dtype=torch.bool)
  254. # sample positive lines
  255. cdx = label.nonzero().flatten()
  256. if len(cdx) > self.n_dyn_posl:
  257. # print("too many positive lines")
  258. perm = torch.randperm(len(cdx), device=device)[: self.n_dyn_posl]
  259. cdx = cdx[perm]
  260. c[cdx] = 1
  261. # sample negative lines
  262. cdx = Lneg[up, vp].nonzero().flatten()
  263. if len(cdx) > self.n_dyn_negl:
  264. # print("too many negative lines")
  265. perm = torch.randperm(len(cdx), device=device)[: self.n_dyn_negl]
  266. cdx = cdx[perm]
  267. c[cdx] = 1
  268. # sample other (unmatched) lines
  269. cdx = torch.randint(len(c), (self.n_dyn_othr,), device=device)
  270. c[cdx] = 1
  271. else:
  272. c = (u < v).flatten()
  273. # sample lines
  274. u, v, label = u[c], v[c], label[c]
  275. xy = xy.reshape(n_type * K, 2)
  276. xyu, xyv = xy[u], xy[v]
  277. u2v = xyu - xyv
  278. u2v /= torch.sqrt((u2v ** 2).sum(-1, keepdim=True)).clamp(min=1e-6)
  279. feat = torch.cat(
  280. [
  281. xyu / 128 * self.use_cood,
  282. xyv / 128 * self.use_cood,
  283. u2v * self.use_slop,
  284. (u[:, None] > K).float(),
  285. (v[:, None] > K).float(),
  286. ],
  287. 1,
  288. )
  289. line = torch.cat([xyu[:, None], xyv[:, None]], 1)
  290. xy = xy.reshape(n_type, K, 2)
  291. jcs = [xy[i, score[i] > 0.03] for i in range(n_type)]
  292. return line, label.float(), feat, jcs
  293. _COMMON_META = {
  294. "categories": _COCO_PERSON_CATEGORIES,
  295. "keypoint_names": _COCO_PERSON_KEYPOINT_NAMES,
  296. "min_size": (1, 1),
  297. }