import cv2 import numpy as np import matplotlib.pyplot as plt import torch import torch.nn.functional as F # from torch.backends import cudnn from fast_reid.fastreid.config import get_cfg from fast_reid.fastreid.modeling.meta_arch import build_model from fast_reid.fastreid.utils.checkpoint import Checkpointer from fast_reid.fastreid.engine import DefaultTrainer, default_argument_parser, default_setup, launch # cudnn.benchmark = True from mmdet.apis import init_detector, inference_detector, show_result_pyplot import mmcv import time def setup_cfg(config_file, opts): # load config from file and command-line arguments cfg = get_cfg() cfg.merge_from_file(config_file) cfg.merge_from_list(opts) cfg.MODEL.BACKBONE.PRETRAIN = False cfg.freeze() return cfg def postprocess(features): # Normalize feature to compute cosine distance features = F.normalize(features) features = features.cpu().data.numpy() return features def preprocess(image, input_size): if len(image.shape) == 3: padded_img = np.ones((input_size[1], input_size[0], 3), dtype=np.uint8) * 114 else: padded_img = np.ones(input_size) * 114 img = np.array(image) r = min(input_size[1] / img.shape[0], input_size[0] / img.shape[1]) resized_img = cv2.resize( img, (int(img.shape[1] * r), int(img.shape[0] * r)), interpolation=cv2.INTER_LINEAR, ) padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img return padded_img, r class FastReIDInterface: def __init__(self, config_file, weights_path, device, batch_size=8): super(FastReIDInterface, self).__init__() if device != 'cpu': self.device = 'cuda' else: self.device = 'cpu' self.batch_size = batch_size self.cfg = setup_cfg(config_file, ['MODEL.WEIGHTS', weights_path]) self.model = build_model(self.cfg) self.model.eval() Checkpointer(self.model).load(weights_path) if self.device != 'cpu': self.model = self.model.eval().to(device='cuda').half() else: self.model = self.model.eval() self.pH, self.pW = self.cfg.INPUT.SIZE_TEST # BoxInstSeg model config_file = '/data/cvprw/AIC23/detect/BoxInstSeg/work_dirs/test-swin-l-640/swin-l-640.py' checkpoint_file = '/data/cvprw/AIC23/detect/BoxInstSeg/work_dirs/test-swin-l-640/iter_10000.pth' self.seg_model = init_detector(config_file, checkpoint_file, device='cuda:0') def inference(self, image, detections): if detections is None or np.size(detections) == 0: return [] H, W, _ = np.shape(image) batch_patches = [] patches_img = [] patches = [] for d in range(np.size(detections, 0)): ltrb = detections[d, :4].astype(np.int_) ltrb[0] = max(0, ltrb[0]) ltrb[1] = max(0, ltrb[1]) ltrb[2] = min(W - 1, ltrb[2]) ltrb[3] = min(H - 1, ltrb[3]) patch_x1 = ltrb[0] patch_y1 = ltrb[1] patch_x2 = ltrb[2] patch_y2 = ltrb[3] # test_size = (100,300) if int(patch_x2-patch_x1)<100 : c_x = int((patch_x2+patch_x1)/2) patch_x1 = max(0, c_x-50) patch_x2 = min(W - 1, c_x+50) if int(patch_y2-patch_y1)<300 : c_y = int((patch_y2+patch_y1)/2) patch_y1 = max(0, c_y-150) patch_y2 = min(H - 1, c_y+150) # print(patch_y1, patch_y2, patch_x1, patch_x2) patch = image[patch_y1:patch_y2, patch_x1:patch_x2, :] # cv2.imwrite("patch.png", patch) result = inference_detector(self.seg_model, patch) img, mask_save = show_result_pyplot(self.seg_model, patch, result, score_thr=0.1) #patch = np.array((ltrb[3] - ltrb[1], ltrb[2] - ltrb[0], 3), dtype=np.uint8) patch = patch[ ltrb[1] - patch_y1 : ltrb[3] - patch_y1 + 1, ltrb[0] - patch_x1 : ltrb[2] - patch_x1 + 1, : ] if mask_save.shape[0] > 0: # person segmented patch_mask = mask_save[ :, ltrb[1] - patch_y1 : ltrb[3] - patch_y1 + 1, ltrb[0] - patch_x1 : ltrb[2] - patch_x1 + 1, ] # find out wich mask to use # normaly detected person's head is tangent to upper bounding box line. channel_idx = patch_mask.sum(axis=(1, 2)).argmax() mask = patch_mask[channel_idx] patch[~mask] = 0 patch = cv2.resize(patch, tuple(self.cfg.INPUT.SIZE_TEST[::-1]), interpolation=cv2.INTER_LINEAR) # patches_img.append(patch) # Make shape with a new batch dimension which is adapted for network input patch = torch.as_tensor(patch.astype("float32").transpose(2, 0, 1)) patch = patch.to(device=self.device).half() patches.append(patch) if (d + 1) % self.batch_size == 0: patches = torch.stack(patches, dim=0) batch_patches.append(patches) patches = [] # save image seg # for idx,patch in enumerate(patches_img): # cv2.imwrite('./seg_save/'+str(frame_id)+'_'+str(idx)+'.png', patch) if len(patches): patches = torch.stack(patches, dim=0) batch_patches.append(patches) features = np.zeros((0, 2048)) # features = np.zeros((0, 768)) for patches in batch_patches: # Run model patches_ = torch.clone(patches) pred = self.model(patches) pred[torch.isinf(pred)] = 1.0 feat = postprocess(pred) nans = np.isnan(np.sum(feat, axis=1)) if np.isnan(feat).any(): for n in range(np.size(nans)): if nans[n]: # patch_np = patches[n, ...].squeeze().transpose(1, 2, 0).cpu().numpy() patch_np = patches_[n, ...] patch_np_ = torch.unsqueeze(patch_np, 0) pred_ = self.model(patch_np_) patch_np = torch.squeeze(patch_np).cpu() patch_np = torch.permute(patch_np, (1, 2, 0)).int() patch_np = patch_np.numpy() plt.figure() plt.imshow(patch_np) plt.show() features = np.vstack((features, feat)) return features