import cv2 import numpy as np from skimage.draw import polygon from skimage.filters import gaussian from skimage.feature import peak_local_max import matplotlib.pyplot as plt import torch import torch.nn.functional as F from cython_nms import nms as cnms from .box_utils import crop, box_iou @torch.no_grad() def draw_lincomb(ids_p, proto_data, masks, img_name, target_dir=None): for kdx in range(masks.shape[0]): target_id = int(ids_p[kdx]) # jdx = kdx + -1 coeffs = masks[kdx, :].cpu().numpy() idx = np.argsort(-np.abs(coeffs)) coeffs_sort = coeffs[idx] arr_h, arr_w = (4, 8) p_h, p_w, _ = proto_data.size() arr_img = np.zeros([p_h * arr_h, p_w * arr_w]) arr_run = np.zeros([p_h * arr_h, p_w * arr_w]) for y in range(arr_h): for x in range(arr_w): i = arr_w * y + x if i == 0: running_total = proto_data[:, :, idx[i]].cpu().numpy() * coeffs_sort[i] else: running_total += proto_data[:, :, idx[i]].cpu().numpy() * coeffs_sort[i] running_total_nonlin = (1 / (1 + np.exp(-running_total))) arr_img[y * p_h:(y + 1) * p_h, x * p_w:(x + 1) * p_w] = (proto_data[:, :, idx[i]] / torch.max( proto_data[:, :, idx[i]])).cpu().numpy() * coeffs_sort[i] arr_run[y * p_h:(y + 1) * p_h, x * p_w:(x + 1) * p_w] = (running_total_nonlin > 0.5).astype(np.float) arr_img = ((arr_img + 1) * 127.5).astype('uint8') arr_img = cv2.applyColorMap(arr_img, cv2.COLORMAP_WINTER) if target_dir is None: cv2.imwrite(f'results/ocid/lincomb_{img_name}', arr_img) else: if not os.path.exists(f'{target_dir}/{target_id}'): os.makedirs(f'{target_dir}/{target_id}') cv2.imwrite(f'{target_dir}/{target_id}/lincomb_{img_name}', arr_img) @torch.no_grad() def fast_nms(cfg, box_pred_kept, cls_pred_kept, ins_coef_pred_kept, grasp_coef_pred_kept): cls_pred_kept, idx = cls_pred_kept.sort(1, descending=True) idx = idx[:, :cfg.top_k] cls_pred_kept = cls_pred_kept[:, :cfg.top_k] num_classes, num_dets = idx.size() box_pred_kept = box_pred_kept[idx.reshape(-1), :].reshape(num_classes, num_dets, 4) ins_coef_pred_kept = ins_coef_pred_kept[idx.reshape(-1), :].reshape(num_classes, num_dets, cfg.num_protos) grasp_coef_pred_kept = grasp_coef_pred_kept[idx.reshape(-1), :].reshape(num_classes, num_dets, 4, cfg.num_protos) # Calculate IoU between predicted boxes iou = box_iou(box_pred_kept, box_pred_kept) iou.triu_(diagonal=1) iou_max, _ = iou.max(dim=1) # Now just filter out the ones higher than the threshold keep = (iou_max <= cfg.nms_iou_thre) # Assign each kept detection to its corresponding class class_ids = torch.arange(num_classes, device=box_pred_kept.device)[:, None].expand_as(keep) class_ids = class_ids[keep] cls_pred_kept = cls_pred_kept[keep] box_pred_kept = box_pred_kept[keep] ins_coef_pred_kept = ins_coef_pred_kept[keep] grasp_coef_pred_kept = grasp_coef_pred_kept[keep] # Only keep the top cfg.max_num_detections highest scores across all classes cls_pred_kept, idx = cls_pred_kept.sort(0, descending=True) idx = idx[:cfg.max_detections] cls_pred_kept = cls_pred_kept[:cfg.max_detections] class_ids = class_ids[idx] box_pred_kept = box_pred_kept[idx] ins_coef_pred_kept = ins_coef_pred_kept[idx] grasp_coef_pred_kept = grasp_coef_pred_kept[idx] return class_ids, cls_pred_kept, box_pred_kept, ins_coef_pred_kept, grasp_coef_pred_kept # IMPORTANT # For batch_size=1 only @torch.no_grad() def ssg_post_processing(cfg, output_dict, data_dict): ori_size = data_dict["ori_size"] ori_h, ori_w = ori_size input_size = max(ori_h, ori_w) protos = output_dict["protos"].squeeze() cls_pred = output_dict["cls_pred"].squeeze() box_pred = output_dict["box_pred"].squeeze() ins_coef_pred = output_dict["ins_coef_pred"].squeeze() grasp_coef_pred = output_dict["grasp_coef_pred"].squeeze() anchors = torch.tensor(output_dict["anchors"], device=protos.device).reshape(-1, 4).squeeze() # print("anchors", anchors.shape) # print("cls_pred", cls_pred.shape) B, N = cls_pred.shape[:2] cls_pred = cls_pred.transpose(1, 0).contiguous() # Exclude the background class cls_pred = cls_pred[1:, :] # get the max score class of 19248 predicted boxes cls_pred_max, _ = torch.max(cls_pred, dim=0) # print("cls_pred_max", cls_pred_max.shape) # filter predicted boxes according the class score keep = (cls_pred_max > cfg.nms_score_thre) # [N_anchors] # print(anchors.shape, cls_pred.shape, box_pred.shape, ins_coef_pred.shape, grasp_coef_pred.shape) # print(keep.shape) anchors_kept = anchors[keep, :] cls_pred_kept = cls_pred[:, keep] box_pred_kept = box_pred[keep, :] ins_coef_pred_kept = ins_coef_pred[keep, :] grasp_coef_pred_kept = grasp_coef_pred[keep, :] # decode boxes box_pred_kept = torch.cat((anchors_kept[:, :2] + box_pred_kept[:, :2] * 0.1 * anchors_kept[:, 2:], anchors_kept[:, 2:] * torch.exp(box_pred_kept[:, 2:] * 0.2)), 1) box_pred_kept[:, :2] -= box_pred_kept[:, 2:] / 2 box_pred_kept[:, 2:] += box_pred_kept[:, :2] box_pred_kept = torch.clip(box_pred_kept, min=0., max=1.) # Fast NMS class_ids, cls_pred_kept, box_pred_kept, ins_coef_pred_kept, grasp_coef_pred_kept = fast_nms(cfg, box_pred_kept, cls_pred_kept, ins_coef_pred_kept, grasp_coef_pred_kept) keep = (cls_pred_kept > 0.3) if not keep.any(): print("No valid instance") else: class_ids = class_ids[keep] cls_pred_kept = cls_pred_kept[keep] box_pred_kept = box_pred_kept[keep] ins_coef_pred_kept = ins_coef_pred_kept[keep] grasp_coef_pred_kept = grasp_coef_pred_kept[keep] class_ids = (class_ids + 1) class_ids = class_ids.cpu().numpy() # if cfg.vis_protos: # ones_coef = torch.ones(pos_coef_p.shape).float().cuda() # # print("ProtoTypes") # draw_lincomb(ids_p, proto_p, ones_coef, "prototypes.png", target_dir) # # print("Semantic") # draw_lincomb(ids_p, proto_p, coef_p, "cogr-sem.png", target_dir) # # print("Grasp pos") # draw_lincomb(ids_p, proto_p, pos_coef_p, "cogr-gr-pos.png", target_dir) # # print("Grasp sin") # draw_lincomb(ids_p, proto_p, sin_coef_p, "cogr-gr-sin.png", target_dir) # # print("Grasp cos") # draw_lincomb(ids_p, proto_p, cos_coef_p, "cogr-gr-cos.png", target_dir) # # print("Grasp wid") # draw_lincomb(ids_p, proto_p, wid_coef_p, "cogr-gr-wid.png", target_dir) ins_masks = torch.sigmoid(torch.matmul(protos, ins_coef_pred_kept.t())).contiguous() grasp_qua_masks = torch.sigmoid(torch.matmul(protos, grasp_coef_pred_kept[:, 0, :].t())).contiguous() grasp_sin_masks = torch.matmul(protos, grasp_coef_pred_kept[:, 1, :].t()).contiguous() grasp_cos_masks = torch.matmul(protos, grasp_coef_pred_kept[:, 2, :].t()).contiguous() grasp_wid_masks = torch.sigmoid(torch.matmul(protos, grasp_coef_pred_kept[:, 3, :].t())).contiguous() ins_masks = crop(ins_masks, box_pred_kept).permute(2,0,1) grasp_qua_masks = crop(grasp_qua_masks, box_pred_kept).permute(2,0,1) grasp_sin_masks = crop(grasp_sin_masks, box_pred_kept).permute(2,0,1) grasp_cos_masks = crop(grasp_cos_masks, box_pred_kept).permute(2,0,1) grasp_wid_masks = crop(grasp_wid_masks, box_pred_kept).permute(2,0,1) ins_masks = F.interpolate(ins_masks.unsqueeze(0), (input_size, input_size), mode='bilinear', align_corners=False).squeeze(0) ins_masks.gt_(0.5) grasp_qua_masks = F.interpolate(grasp_qua_masks.unsqueeze(0), (input_size, input_size), mode='bilinear', align_corners=False).squeeze(0) grasp_sin_masks = F.interpolate(grasp_sin_masks.unsqueeze(0), (input_size, input_size), mode='bilinear', align_corners=False).squeeze(0) grasp_cos_masks = F.interpolate(grasp_cos_masks.unsqueeze(0), (input_size, input_size), mode='bilinear', align_corners=False).squeeze(0) grasp_wid_masks = F.interpolate(grasp_wid_masks.unsqueeze(0), (input_size, input_size), mode='bilinear', align_corners=False).squeeze(0) ins_masks = ins_masks[:, 0:ori_h, 0:ori_w].cpu().numpy() grasp_qua_masks = grasp_qua_masks[:, 0:ori_h, 0:ori_w].cpu().numpy() grasp_sin_masks = grasp_sin_masks[:, 0:ori_h, 0:ori_w].cpu().numpy() grasp_cos_masks = grasp_cos_masks[:, 0:ori_h, 0:ori_w].cpu().numpy() grasp_wid_masks = grasp_wid_masks[:, 0:ori_h, 0:ori_w].cpu().numpy() grasp_ang_masks = [] for i in range(ins_masks.shape[0]): grasp_qua_masks[i] = gaussian(grasp_qua_masks[i], 2.0, preserve_range=True) ang_mask = (np.arctan2(grasp_sin_masks[i], grasp_cos_masks[i]) / 2.0) grasp_ang_masks.append(ang_mask) grasp_ang_masks = np.asarray(grasp_ang_masks) scale = np.array([ori_w, ori_w, ori_w, ori_w]) box_pred_kept = box_pred_kept.cpu().numpy() * scale ins_grasp_rects_top1 = [] ins_grasp_rects_top5 = [] for i in range(ins_masks.shape[0]): grasps_top1, _ = detect_grasps(grasp_qua_masks[i], grasp_sin_masks[i], grasp_cos_masks[i], grasp_wid_masks[i], 1) grasps_top5, _ = detect_grasps(grasp_qua_masks[i], grasp_sin_masks[i], grasp_cos_masks[i], grasp_wid_masks[i], 5) ins_grasp_rects_top1.append(grasps_top1) ins_grasp_rects_top5.append(grasps_top5) return { "cls": class_ids, "bboxes": box_pred_kept, "ins_masks": ins_masks, "grasps_top1": ins_grasp_rects_top1, "grasps_top5": ins_grasp_rects_top5, "grasp_masks": (grasp_qua_masks, grasp_ang_masks, grasp_wid_masks) } # def visualization(img, mask, grasp_masks, grasps, text, save_path=None): # grasp_qua_mask, grasp_ang_mask, grasp_wid_mask = grasp_masks # fig = plt.figure(figsize=(25, 10)) # # draw grasp rectangles in image # tmp = img.copy() # for rect in grasps: # center_x, center_y, width, height, theta = rect # box = ((center_x, center_y), (width, height), -(theta+180)) # box = cv2.boxPoints(box) # box = np.intp(box) # ptA, ptB, ptC, ptD = [list(map(int, pt.tolist())) for pt in box] # tmp = cv2.line(tmp, ptA, ptB, (0,0,0xff), 2) # tmp = cv2.line(tmp, ptD, ptC, (0,0,0xff), 2) # tmp = cv2.line(tmp, ptB, ptC, (0xff,0,0), 2) # tmp = cv2.line(tmp, ptA, ptD, (0xff,0,0), 2) # # tmp = cv2.rectangle(tmp, (bbox[0],bbox[1]), (bbox[2],bbox[3]), (0,255,0), 2) # tmp = cv2.putText(tmp, text, (0,10), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,0), 2, cv2.LINE_AA) # # draw predicted instance mask in image # msk_img = (img * 0.3).astype(np.uint8).copy() # mask = mask.astype(np.uint8) # msk_img[mask, 0] = 255 # ax = fig.add_subplot(2, 3, 1) # ax.imshow(img/255.) # ax.set_title('RGB') # ax.axis('off') # ax = fig.add_subplot(2, 3, 2) # ax.imshow(tmp/255.) # ax.set_title('predicted grasps') # ax.axis('off') # ax = fig.add_subplot(2, 3, 3) # ax.imshow(mask) # ax.set_title('predicted instance mask') # ax.axis('off') # ax = fig.add_subplot(2, 3, 4) # plot = ax.imshow(grasp_qua_mask, cmap='jet', vmin=0, vmax=1) # ax.set_title('Grasp quality') # ax.axis('off') # plt.colorbar(plot) # ax = fig.add_subplot(2, 3, 5) # plot = ax.imshow(grasp_ang_mask, cmap='jet', vmin=0, vmax=1) # ax.set_title('Grasp Angle') # ax.axis('off') # plt.colorbar(plot) # ax = fig.add_subplot(2, 3, 6) # plot = ax.imshow(grasp_wid_mask, cmap='jet', vmin=0, vmax=1) # ax.set_title('Grasp Width') # ax.axis('off') # plt.colorbar(plot) # plt.suptitle(f"{text}", fontsize=20) # plt.tight_layout() # plt.savefig(save_path) def visualization(img, mask, grasp_masks, grasps, text, save_path=None): """ img: RGB uint8 (H, W, 3) mask: None or (H, W) binary/float/uint8 grasp_masks: (grasp_qua_mask, grasp_ang_mask, grasp_wid_mask) grasps: list of rectangles """ grasp_qua_mask, grasp_ang_mask, grasp_wid_mask = grasp_masks # --- draw grasp rectangles in image --- fig = plt.figure(figsize=(25, 10)) tmp = img.copy() for rect in grasps: center_x, center_y, width, height, theta = rect box = ((center_x, center_y), (width, height), -(theta + 180)) box = cv2.boxPoints(box) box = np.intp(box) ptA, ptB, ptC, ptD = [list(map(int, pt.tolist())) for pt in box] tmp = cv2.line(tmp, ptA, ptB, (0, 0, 0xFF), 2) tmp = cv2.line(tmp, ptD, ptC, (0, 0, 0xFF), 2) tmp = cv2.line(tmp, ptB, ptC, (0xFF, 0, 0), 2) tmp = cv2.line(tmp, ptA, ptD, (0xFF, 0, 0), 2) tmp = cv2.putText( tmp, text, (0, 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2, cv2.LINE_AA ) has_mask = mask is not None # --- optional instance mask overlay --- if has_mask: # Accept bool / float / uint8 mask_np = np.asarray(mask) if mask_np.dtype != np.bool_: # if it's float/prob, consider >0.5 as foreground if np.issubdtype(mask_np.dtype, np.floating): mask_np = mask_np > 0.5 else: mask_np = mask_np.astype(np.uint8) > 0 msk_img = (img * 0.3).astype(np.uint8).copy() msk_img[mask_np, 0] = 255 # highlight in red channel # --- layout --- # With mask: 2x3 like before # Without mask: 2x2 (remove the mask subplot) if has_mask: ax = fig.add_subplot(2, 3, 1) ax.imshow(img / 255.) ax.set_title('RGB') ax.axis('off') ax = fig.add_subplot(2, 3, 2) ax.imshow(tmp / 255.) ax.set_title('predicted grasps') ax.axis('off') ax = fig.add_subplot(2, 3, 3) ax.imshow(mask_np.astype(np.uint8)) ax.set_title('predicted instance mask') ax.axis('off') ax = fig.add_subplot(2, 3, 4) plot = ax.imshow(grasp_qua_mask, cmap='jet', vmin=0, vmax=1) ax.set_title('Grasp quality') ax.axis('off') plt.colorbar(plot) ax = fig.add_subplot(2, 3, 5) plot = ax.imshow(grasp_ang_mask, cmap='jet', vmin=0, vmax=1) ax.set_title('Grasp Angle') ax.axis('off') plt.colorbar(plot) ax = fig.add_subplot(2, 3, 6) plot = ax.imshow(grasp_wid_mask, cmap='jet', vmin=0, vmax=1) ax.set_title('Grasp Width') ax.axis('off') plt.colorbar(plot) else: ax = fig.add_subplot(2, 2, 1) ax.imshow(img / 255.) ax.set_title('RGB') ax.axis('off') ax = fig.add_subplot(2, 2, 2) ax.imshow(tmp / 255.) ax.set_title('predicted grasps') ax.axis('off') ax = fig.add_subplot(2, 2, 3) plot = ax.imshow(grasp_qua_mask, cmap='jet', vmin=0, vmax=1) ax.set_title('Grasp quality') ax.axis('off') plt.colorbar(plot) ax = fig.add_subplot(2, 2, 4) plot = ax.imshow(grasp_ang_mask, cmap='jet', vmin=0, vmax=1) ax.set_title('Grasp Angle') ax.axis('off') plt.colorbar(plot) plt.suptitle(f"{text}", fontsize=20) plt.tight_layout() if save_path is not None: plt.savefig(save_path) plt.close(fig) def detect_grasps(grasp_quality_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask, num_grasps=5): grasps = [] max_width = 100 local_max = peak_local_max(grasp_quality_mask, min_distance=2, threshold_abs=0.4, num_peaks=num_grasps) grasp_angle_mask = (np.arctan2(grasp_sin_mask, grasp_cos_mask) / 2.0) for p_array in local_max: grasp_point = tuple(p_array) grasp_angle = grasp_angle_mask[grasp_point] / np.pi * 180 grasp_width = grasp_wid_mask[grasp_point] grasps.append([float(grasp_point[1]), float(grasp_point[0]), grasp_width*max_width, 20, grasp_angle]) return grasps, grasp_angle_mask def calculate_iou(rect_p, rect_gt, shape=(480, 640), angle_threshold=30): if abs(rect_p[4] - rect_gt[4]) > angle_threshold and abs(rect_p[4] + rect_gt[4]) > angle_threshold: return 0 center_x, center_y, w_rect, h_rect, theta, _ = rect_gt gt_r_rect = ((center_x, center_y), (w_rect, h_rect), -theta) gt_box = cv2.boxPoints(gt_r_rect) gt_box = np.int0(gt_box) rr1, cc1 = polygon(gt_box[:, 0], gt_box[:,1], shape) mask_rr = rr1 < shape[1] rr1 = rr1[mask_rr] cc1 = cc1[mask_rr] mask_cc = cc1 < shape[0] cc1 = cc1[mask_cc] rr1 = rr1[mask_cc] center_x, center_y, w_rect, h_rect, theta = rect_p p_r_rect = ((center_x, center_y), (w_rect, h_rect), -theta) p_box = cv2.boxPoints(p_r_rect) p_box = np.int0(p_box) rr2, cc2 = polygon(p_box[:, 0], p_box[:,1], shape) mask_rr = rr2 < shape[1] rr2 = rr2[mask_rr] cc2 = cc2[mask_rr] mask_cc = cc2 < shape[0] cc2 = cc2[mask_cc] rr2 = rr2[mask_cc] area = np.zeros(shape) area[cc1, rr1] += 1 area[cc2, rr2] += 1 union = np.sum(area > 0) intersection = np.sum(area == 2) if union <= 0: return 0 else: return intersection / union def calculate_max_iou(rects_p, rects_gt): max_iou = 0 for rect_gt in rects_gt: for rect_p in rects_p: iou = calculate_iou(rect_p, rect_gt) # print("==============================") # print(rect_p, rect_gt, iou) if iou > max_iou: max_iou = iou return max_iou def calculate_jacquard_index(grasp_preds, grasp_targets, iou_threshold=0.25): j_index = 0 grasp_preds = np.asarray(grasp_preds) grasp_targets = np.asarray(grasp_targets) grasp_targets[:, 3] = 20 grasp_targets[:, 2] = np.clip(grasp_targets[:, 2], 0, 100) iou = calculate_max_iou(grasp_preds, grasp_targets) if iou > iou_threshold: j_index = 1 return j_index