{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Objective\n\nIn this notebook, you'll see how to use YoloV7 and detectron2 for players instance segmentation. \n\nI created the notebook following this amazing tutorial [Instance Segmentation with YOLOv7 in Python](https://stackabuse.com/instance-segmentation-with-yolov7-in-python/) on how do instance semgentation using, \n\nI didn't write <s>any</s> all of the code from scratch, it belongs to David Landup who is the author of the tutorial.\n\n#### Update: 09-01-2023 \n\nI have added a added a snippet in the plot function that matches each of the predicted bboxes with helmets from the frame.","metadata":{}},{"cell_type":"markdown","source":"### Extract one video to test","metadata":{}},{"cell_type":"code","source":"!mkdir -p frames\n!ffmpeg -i /kaggle/input/nfl-player-contact-detection/train/58168_003392_Sideline.mp4 -q:v 2 -f image2 /kaggle/working/frames/frame_%04d.jpg -hide_banner -loglevel error","metadata":{"execution":{"iopub.status.busy":"2023-01-09T10:25:43.336644Z","iopub.execute_input":"2023-01-09T10:25:43.337073Z","iopub.status.idle":"2023-01-09T10:25:53.335743Z","shell.execute_reply.started":"2023-01-09T10:25:43.336984Z","shell.execute_reply":"2023-01-09T10:25:53.334131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Download Yolov7 and Detectron2","metadata":{}},{"cell_type":"code","source":"! pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html\n\n! git clone -b mask https://github.com/WongKinYiu/yolov7.git\n! pip install pyyaml==5.1\n! pip install 'git+https://github.com/facebookresearch/detectron2.git'\n\n%cd yolov7\n! curl -L https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-mask.pt -o yolov7-mask.pt","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-01-09T10:25:53.342221Z","iopub.execute_input":"2023-01-09T10:25:53.344462Z","iopub.status.idle":"2023-01-09T10:32:42.232647Z","shell.execute_reply.started":"2023-01-09T10:25:53.344421Z","shell.execute_reply":"2023-01-09T10:32:42.231451Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\nimport cv2\nimport yaml\nfrom torchvision import transforms\nimport numpy as np\n\nfrom glob import glob\n\nfrom utils.datasets import letterbox\nfrom utils.general import non_max_suppression_mask_conf\n\nfrom detectron2.modeling.poolers import ROIPooler\nfrom detectron2.structures import Boxes\nfrom detectron2.utils.memory import retry_if_cuda_oom\nfrom detectron2.layers import paste_masks_in_image","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:14.23407Z","iopub.execute_input":"2023-01-09T11:12:14.234465Z","iopub.status.idle":"2023-01-09T11:12:14.240659Z","shell.execute_reply.started":"2023-01-09T11:12:14.234432Z","shell.execute_reply":"2023-01-09T11:12:14.239662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"frames_paths = sorted(glob('/kaggle/working/frames/*'))\n\ni = 300\nframe = cv2.imread(frames_paths[i])\nframe = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n\nfig = plt.figure(figsize=(12, 6))\nplt.imshow(frame);","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:14.472882Z","iopub.execute_input":"2023-01-09T11:12:14.473254Z","iopub.status.idle":"2023-01-09T11:12:14.977589Z","shell.execute_reply.started":"2023-01-09T11:12:14.473223Z","shell.execute_reply":"2023-01-09T11:12:14.976643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls /kaggle/working/frames/ | wc -l","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:14.979521Z","iopub.execute_input":"2023-01-09T11:12:14.98038Z","iopub.status.idle":"2023-01-09T11:12:16.011426Z","shell.execute_reply.started":"2023-01-09T11:12:14.980338Z","shell.execute_reply":"2023-01-09T11:12:16.010161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\ntr_helmets = pd.read_csv('/kaggle/input/nfl-player-contact-detection/train_baseline_helmets.csv')\ntr_helmets.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:16.014158Z","iopub.execute_input":"2023-01-09T11:12:16.014609Z","iopub.status.idle":"2023-01-09T11:12:19.549319Z","shell.execute_reply.started":"2023-01-09T11:12:16.014561Z","shell.execute_reply":"2023-01-09T11:12:19.548372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"frames_paths[i]","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:19.552073Z","iopub.execute_input":"2023-01-09T11:12:19.552465Z","iopub.status.idle":"2023-01-09T11:12:19.558643Z","shell.execute_reply.started":"2023-01-09T11:12:19.552426Z","shell.execute_reply":"2023-01-09T11:12:19.557617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"frame_helmets = tr_helmets.query('video == \"58168_003392_Sideline.mp4\" and frame == 301')\nframe_helmets.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:19.560324Z","iopub.execute_input":"2023-01-09T11:12:19.56109Z","iopub.status.idle":"2023-01-09T11:12:19.687004Z","shell.execute_reply.started":"2023-01-09T11:12:19.561052Z","shell.execute_reply":"2023-01-09T11:12:19.68605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"frame_helmets.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:19.689608Z","iopub.execute_input":"2023-01-09T11:12:19.69032Z","iopub.status.idle":"2023-01-09T11:12:19.705567Z","shell.execute_reply.started":"2023-01-09T11:12:19.69028Z","shell.execute_reply":"2023-01-09T11:12:19.704524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_img = frame.copy()\n\nfor helmet in frame_helmets.itertuples():\n    startpoint = (helmet.left, helmet.top)\n    endpoint = (helmet.left + helmet.width, helmet.top + helmet.height)\n    cv2.rectangle(pred_img, startpoint, endpoint, (0, 0, 0), 1)\n    \n\nplt.figure(figsize=(15, 15))\nplt.imshow(pred_img);","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:19.707224Z","iopub.execute_input":"2023-01-09T11:12:19.70761Z","iopub.status.idle":"2023-01-09T11:12:20.442278Z","shell.execute_reply.started":"2023-01-09T11:12:19.707574Z","shell.execute_reply":"2023-01-09T11:12:20.441395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom scipy.optimize import linear_sum_assignment\n\ndef match_bboxes_with_helmets(frame_helmets, nbboxes):\n    helmets_coords = frame_helmets[['left', 'top']].values\n    helmets_coords[:, 0] = helmets_coords[:, 0] * (640 / 1280)\n    helmets_coords[:, 1] = helmets_coords[:, 1] * (384 / 720)\n\n    # Calculate the cost matrix\n    cost_matrix = np.zeros((len(nbboxes), len(helmets_coords)))\n    for i in range(len(nbboxes)):\n        cost_matrix[i, :] = np.linalg.norm(nbboxes[i, [0,1]] - helmets_coords, axis=1)\n\n    # Use the Hungarian algorithm to find the optimal match\n    row_ind, col_ind = linear_sum_assignment(cost_matrix)\n    \n    # The indices in row_ind and col_ind give the indices of the matched coordinates\n#     matched_bboxes = [nbboxes[j] for j in row_ind]\n    matched_helmets = [frame_helmets.iloc[i] for i in col_ind]\n    \n    return matched_helmets\n\n\ndef load_model():\n    model = torch.load('yolov7-mask.pt', map_location=device)['model']\n    # Put in inference mode\n    model.eval()\n\n    if torch.cuda.is_available():\n        # half() turns predictions into float16 tensors\n        # which significantly lowers inference time\n        model.half().to(device)\n    return model\n\ndef run_inference(url):\n    image = cv2.imread(url) # shape: (480, 640, 3)\n    # Resize and pad image\n    image = letterbox(image, 640, stride=64, auto=True)[0] # shape: (480, 640, 3)\n    # Apply transforms\n    image = transforms.ToTensor()(image) # torch.Size([3, 480, 640])\n    # Match tensor type (`torch.FloatTensor` -> `torch.HalfTensor`) with model\n    image = image.half().to(device)\n    # Turn image into batch\n    image = image.unsqueeze(0) # torch.Size([1, 3, 480, 640])\n    output = model(image)\n    return output, image\n\n\ndef plot_results(original_image, pred_img, pred_masks_np, nbboxes, pred_cls, pred_conf, frame_helmets, plot_labels=True):\n    \n\n    matched_helmets = match_bboxes_with_helmets(frame_helmets, nbboxes)\n    \n    \n    for i, (one_mask, bbox, cls, conf) in enumerate(zip(pred_masks_np, nbboxes, pred_cls, pred_conf)):\n        if conf < 0.1:\n            continue\n        color = [np.random.randint(255), np.random.randint(255), np.random.randint(255)]\n\n        pred_img = pred_img.copy()\n\n        # Apply mask over image in color\n        pred_img[one_mask] = pred_img[one_mask] * 0.5 + np.array(color, dtype=np.uint8) * 0.5\n        # Draw rectangles around all found objects\n        pred_img = cv2.rectangle(pred_img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 1)\n        \n        helmet = matched_helmets[i]\n        \n        startpoint = (int(helmet.left * (640 / 1280)), int(helmet.top * (384 / 720)))\n        endpoint = (int((helmet.left + helmet.width) * (640 / 1280)), int((helmet.top + helmet.height) * (384 / 720)))\n        cv2.rectangle(pred_img, startpoint, endpoint, color, 1)\n\n        if plot_labels:\n            label = '%s %.3f' % (names[int(cls)], conf)\n            t_size = cv2.getTextSize(label, 0, fontScale=0.1, thickness=1)[0]\n            c2 = bbox[0] + t_size[0], bbox[1] - t_size[1] - 3\n            pred_img = cv2.rectangle(pred_img, (bbox[0], bbox[1]), c2, color, -1, cv2.LINE_AA)\n            pred_img = cv2.putText(pred_img, label, (bbox[0], bbox[1] - 2), 0, 0.5, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA)  \n                    \n\n    fig, ax = plt.subplots(1, 2, figsize=(pred_img.shape[0]/10, pred_img.shape[1]/10), dpi=150)\n\n    original_image = np.moveaxis(image.cpu().numpy().squeeze(), 0, 2).astype('float32')\n    original_image = cv2.cvtColor(original_image, cv2.COLOR_RGB2BGR)\n\n    ax[0].imshow(original_image)\n    ax[0].axis(\"off\")\n    ax[1].imshow(pred_img)\n    ax[1].axis(\"off\")","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-01-09T11:12:20.444137Z","iopub.execute_input":"2023-01-09T11:12:20.444525Z","iopub.status.idle":"2023-01-09T11:12:20.472668Z","shell.execute_reply.started":"2023-01-09T11:12:20.444486Z","shell.execute_reply":"2023-01-09T11:12:20.471691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open('data/hyp.scratch.mask.yaml') as f:\n    hyp = yaml.load(f, Loader=yaml.FullLoader)\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = load_model()","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:20.477076Z","iopub.execute_input":"2023-01-09T11:12:20.477344Z","iopub.status.idle":"2023-01-09T11:12:20.675957Z","shell.execute_reply.started":"2023-01-09T11:12:20.477319Z","shell.execute_reply":"2023-01-09T11:12:20.674955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output, image = run_inference(frames_paths[i])","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-01-09T11:12:20.678114Z","iopub.execute_input":"2023-01-09T11:12:20.67849Z","iopub.status.idle":"2023-01-09T11:12:20.731711Z","shell.execute_reply.started":"2023-01-09T11:12:20.678444Z","shell.execute_reply":"2023-01-09T11:12:20.730611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output.keys()","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-01-09T11:12:20.733372Z","iopub.execute_input":"2023-01-09T11:12:20.733759Z","iopub.status.idle":"2023-01-09T11:12:20.740354Z","shell.execute_reply.started":"2023-01-09T11:12:20.733723Z","shell.execute_reply":"2023-01-09T11:12:20.739334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"inf_out = output['test']\nattn = output['attn']\nbases = output['bases']\nsem_output = output['sem']\n\nbases = torch.cat([bases, sem_output], dim=1)\nnb, _, height, width = image.shape\nnames = model.names\npooler_scale = model.pooler_scale\n\npooler = ROIPooler(output_size=hyp['mask_resolution'], \n                   scales=(pooler_scale,), \n                   sampling_ratio=1, \n                   pooler_type='ROIAlignV2', \n                   canonical_level=2)\n                   \n# output, output_mask, output_mask_score, output_ac, output_ab\noutput, output_mask, _, _, _ = non_max_suppression_mask_conf(inf_out, \n                                                             attn, \n                                                             bases, \n                                                             pooler, \n                                                             hyp, \n                                                             conf_thres=0.25, \n                                                             iou_thres=0.65, \n                                                             merge=False, \n                                                             mask_iou=None)","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:20.741643Z","iopub.execute_input":"2023-01-09T11:12:20.742585Z","iopub.status.idle":"2023-01-09T11:12:20.756311Z","shell.execute_reply.started":"2023-01-09T11:12:20.74253Z","shell.execute_reply":"2023-01-09T11:12:20.755315Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output[0].shape \noutput_mask[0].shape","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-01-09T11:12:20.757587Z","iopub.execute_input":"2023-01-09T11:12:20.758079Z","iopub.status.idle":"2023-01-09T11:12:20.76835Z","shell.execute_reply.started":"2023-01-09T11:12:20.758041Z","shell.execute_reply":"2023-01-09T11:12:20.767446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred, pred_masks = output[0], output_mask[0]\nbase = bases[0]\nbboxes = Boxes(pred[:, :4])\n\noriginal_pred_masks = pred_masks.view(-1, \n                                      hyp['mask_resolution'], \n                                      hyp['mask_resolution'])\n\npred_masks = retry_if_cuda_oom(paste_masks_in_image)(original_pred_masks, \n                                                     bboxes, \n                                                     (height, width), \n                                                     threshold=0.5)\n                                                     \n# Detach Tensors from the device, send to the CPU and turn into NumPy arrays\npred_masks_np = pred_masks.detach().cpu().numpy()\npred_cls = pred[:, 5].detach().cpu().numpy()\npred_conf = pred[:, 4].detach().cpu().numpy()\nnimg = image[0].permute(1, 2, 0) * 255\nnimg = nimg.cpu().numpy().astype(np.uint8)\nnimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)\nnbboxes = bboxes.tensor.detach().cpu().numpy().astype(np.int32)","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:20.770048Z","iopub.execute_input":"2023-01-09T11:12:20.770692Z","iopub.status.idle":"2023-01-09T11:12:20.79041Z","shell.execute_reply.started":"2023-01-09T11:12:20.770656Z","shell.execute_reply":"2023-01-09T11:12:20.789391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline\nplot_results(image, nimg, pred_masks_np, nbboxes, pred_cls, pred_conf, frame_helmets, plot_labels=False)","metadata":{"execution":{"iopub.status.busy":"2023-01-09T11:12:20.792027Z","iopub.execute_input":"2023-01-09T11:12:20.792628Z","iopub.status.idle":"2023-01-09T11:12:22.24103Z","shell.execute_reply.started":"2023-01-09T11:12:20.792592Z","shell.execute_reply":"2023-01-09T11:12:22.239325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## How can this be used?\n\n1. Off the top my head, I think this can be paired with helmets to identify the masks for each of the players.\n2. This can be used to identify the pair of players whose contact is in question in a single cropped frame.\n2. It could also be used for background removal and isolation of the pair of players.\n3. It can also remove the need of cropping at all, since we can feed the data to a model and output the contact between every possible pair of players.\n\n## I hope you enjoyed.\n\n## Don't forget to upvote if you find this useful","metadata":{}}]}