{"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":"code","source":"!pip install ultralytics -q\n!pip install pycocotools\n\nfrom ultralytics import YOLO\nimport pycocotools\nfrom pycocotools.coco import COCO\n\nimport os\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport json\nimport pandas as pd","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def binarize(img):\n#     gim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n#     gim = cv2.adaptiveThreshold(gim, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 45, 11)\n#     g3im = cv2.cvtColor(gim, cv2.COLOR_GRAY2BGR)\n#     return g3im","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nfrom tqdm.notebook import tqdm\nimport warnings\nwarnings.filterwarnings(\"ignore\")\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# yolo_model = YOLO('/kaggle/input/yolo-runs5/best (1).pt')\nyolo_model = YOLO('/kaggle/input/weights/best.pt')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_metadata = '/kaggle/input/dlsprint2/badlad/badlad-test-metadata.json'\ntest_image_path = '/kaggle/input/dlsprint2/badlad/images/test'\ntest_coco = COCO(test_metadata)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train_label = '/kaggle/input/dlsprint2/badlad/labels/coco_format/train/badlad-train-coco.json'\n# train_image_path = '/kaggle/input/dlsprint2/badlad/images/train'\n# train_coco = COCO(train_label)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_encode(mask: np.ndarray) -> str:\n    pixels = mask.T.flatten()\n\n    use_padding = False\n    if pixels[0] or pixels[-1]:\n        use_padding = True\n        pixel_padded = np.zeros([len(pixels) + 2], dtype=pixels.dtype)\n        pixel_padded[1:-1] = pixels\n        pixels = pixel_padded\n\n    rle = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    if use_padding:\n        rle = rle - 1  \n    rle[1::2] = rle[1::2] - rle[:-1:2]\n    return ' '.join(str(x) for x in rle)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def yolo_predicted_mask(result, ht, wt):\n    final_mask = {0:np.zeros((ht, wt)), \n                  1:np.zeros((ht, wt)), \n                  2:np.zeros((ht, wt)),\n                  3:np.zeros((ht, wt))}\n    \n    for idx, obj in enumerate(result):\n        label = obj['class']\n        if len(obj['segments']['x']) != 0 and label != 2 and label != 3:\n            points = np.array(list(zip(obj['segments']['x'], obj['segments']['y'])), dtype=np.int32)\n            final_mask[label] = cv2.fillPoly(final_mask[label], pts=[points], color=1)\n        else:\n            x1 = obj['box']['x1']\n            y1 = obj['box']['y1']\n            x2 = obj['box']['x2']\n            y2 = obj['box']['y2']\n            points = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2], [x1, y1]], dtype=np.int32)\n        final_mask[label] = cv2.fillPoly(final_mask[label], pts=[points], color=1)\n    return final_mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def yolo_predicted_mask1(Result, ht, wt):\n#     final_mask = {0:np.zeros((ht, wt)),\n#                   1:np.zeros((ht, wt)),\n#                   2:np.zeros((ht, wt)),\n#                   3:np.zeros((ht, wt))}\n#     for r in Result:\n#         cat_id = r.boxes.cls\n#         if len(cat_id) > 0:\n#             for index, points in enumerate(r.masks.xy):\n#                 pts = points.astype(np.int32)\n#                 label = int(cat_id[index])\n#                 final_mask[label] = cv2.fillPoly(final_mask[label], pts=[pts], color=1)\n#     return final_mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sliding_window(image, kernel_height, y_stride):\n    y_start, y_end = 0, kernel_height\n    isheight = 1\n    while isheight:\n        yield image[y_start:y_end], y_start, y_end\n        if y_end == image.shape[0]:\n            isheight = 0\n        y_start += y_stride\n        y_end = min(y_start+kernel_height, image.shape[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Id, Prediction = [], []\nIds = test_coco.getImgIds()\nfor x in Ids:\n    print(x)\n    image_dict = test_coco.loadImgs(ids=x)\n    image_path = os.path.join(test_image_path, image_dict[0]['file_name'])\n    test_image = cv2.imread(image_path)\n    height, width = test_image.shape[0], test_image.shape[1]\n    \n    if test_image.shape[0] > 2000 and test_image.shape[1] > 2000:\n        Mask = {0:np.zeros((height, width)), \n              1:np.zeros((height, width)), \n              2:np.zeros((height, width)),\n              3:np.zeros((height, width))}\n\n        k_ht = height//3\n        ystride = k_ht - int(k_ht*0.3)\n    \n        for slc, ystart, yend in sliding_window(test_image, k_ht, ystride):\n            result = yolo_model.predict(slc, device=device, imgsz=1024, conf=0.4)\n            result = result[0].tojson()\n            result = json.loads(result)\n            mask = yolo_predicted_mask(result, slc.shape[0], slc.shape[1])\n            Mask[0][ystart:yend] = mask[0]\n            Mask[1][ystart:yend] = mask[1]\n            Mask[2][ystart:yend] = mask[2]\n            Mask[3][ystart:yend] = mask[3]\n    else:\n        result = yolo_model.predict(test_image, device=device, imgsz=1024, conf=0.4)\n        result = result[0].tojson()\n        result = json.loads(result)\n        Mask = yolo_predicted_mask(result, height, width)\n        \n    for k in Mask.keys():\n        rle = rle_encode(Mask[k])\n        Id.append(str(image_dict[0]['id'])+'_'+str(k))\n        Prediction.append(rle)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Id, Prediction = [], []\n# Ids = test_coco.getImgIds()\n# for x in Ids:\n#     image_dict = test_coco.loadImgs(ids=x)\n#     image_path = os.path.join(test_image_path, image_dict[0]['file_name'])\n#     test_image = cv2.imread(image_path)\n#     results = yolo_model.predict(test_image, device=device, imgsz=1024)\n#     results = results[0].tojson()\n#     results = json.loads(results)\n#     mask = yolo_predicted_mask(results, test_image.shape[0], test_image.shape[1])\n#     for k in mask.keys():\n#         rle = rle_encode(mask[k])\n#         Id.append(str(image_dict[0]['id'])+'_'+str(k))\n#         Prediction.append(rle)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame({'Id':Id, 'Predicted':Prediction}, index=None)\nsubmission.to_csv('submission.csv', index=None)\n# submission","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Ids = train_coco.getImgIds(imgIds=[7812])\n# Ids\n# submission","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# image_dict = train_coco.loadImgs(ids=Ids)\n# image_path = os.path.join(train_image_path, image_dict[38]['file_name'])\n# test_image = cv2.imread(image_path)\n# # test_image = binarize(test_image)\n# height, width = test_image.shape[0], test_image.shape[1]\n\n# Mask = {0:np.zeros((height, width)), \n#       1:np.zeros((height, width)), \n#       2:np.zeros((height, width)),\n#       3:np.zeros((height, width))}\n\n# k_ht = test_image.shape[0]//3\n# stride = k_ht - int(k_ht*0.3)\n# # imgsz=(slc.shape[0], slc.shape[1])\n# for slc, start, end in sliding_window(test_image, k_ht, stride):\n#     result = yolo_model.predict(slc, device=device, imgsz=1024)\n#     result = result[0].tojson()\n#     result = json.loads(result)\n#     mask = yolo_predicted_mask(result, slc.shape[0], slc.shape[1])\n#     Mask[0][start:end] = mask[0]\n#     Mask[1][start:end] = mask[1]\n#     Mask[2][start:end] = mask[2]\n#     Mask[3][start:end] = mask[3]\n\n# plt.figure(figsize=(12, 10))\n# plt.subplot(2,3,1)\n# plt.imshow(test_image)\n\n# plt.subplot(2,3,2)\n# plt.imshow(Mask[0])\n\n# plt.subplot(2,3,3)\n# plt.imshow(Mask[1])\n\n# plt.subplot(2,3,4)\n# plt.imshow(Mask[2])\n\n# plt.subplot(2,3,5)\n# plt.imshow(Mask[3])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# visualize_image(Ids[28], train_coco, train_image_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# image_dict = train_coco.loadImgs(ids=Ids)\n# image_path = os.path.join(train_image_path, image_dict[28]['file_name'])\n# test_image = cv2.imread(image_path)\n# test_image = binarize(test_image)\n# height, width = test_image.shape[0], test_image.shape[1]\n# results = yolo_model.predict(test_image, device=device)\n# results = results[0].tojson()\n# results = json.loads(results)\n# mask = yolo_predicted_mask(results, height, width)\n\n# plt.figure(figsize=(12, 10))\n# plt.subplot(2,3,1)\n# plt.imshow(test_image)\n\n# plt.subplot(2,3,2)\n# plt.imshow(mask[0])\n\n# plt.subplot(2,3,3)\n# plt.imshow(mask[1])\n\n# plt.subplot(2,3,4)\n# plt.imshow(mask[2])\n\n# plt.subplot(2,3,5)\n# plt.imshow(mask[3])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def visualize_image(index, coco, path):\n# #     img_ids = coco.getImgIds(imgIds=index)\n#     images = coco.loadImgs(index)\n#     plt.figure(figsize=(15, 12))\n\n#     plt.subplot(2,3,1)\n#     normal_image = os.path.join(path, images[0]['file_name'])\n#     normal_image = cv2.imread(normal_image)\n#     height, width = normal_image.shape[0], normal_image.shape[1]\n#     plt.imshow(normal_image)\n#     plt.title('original image')\n    \n#     masks = {0:np.zeros((height, width)), \n#       1:np.zeros((height, width)), \n#       2:np.zeros((height, width)),\n#       3:np.zeros((height, width))}\n    \n#     annot_id = coco.getAnnIds(imgIds=index)\n#     annotations = coco.loadAnns(annot_id)\n#     for ann in annotations:\n#         label = ann['category_id']\n#         m = coco.annToMask(ann)\n#         masks[label] += m\n    \n#     plt.subplot(2,3,2)\n#     plt.imshow(masks[0])\n#     plt.title('Paragraph')\n\n#     plt.subplot(2,3,3)\n#     plt.imshow(masks[1])\n#     plt.title('text_box')\n\n#     plt.subplot(2,3,4)\n#     plt.imshow(masks[2])\n#     plt.title('image')\n    \n#     plt.subplot(2,3,5)\n#     plt.imshow(masks[3])\n#     plt.title('table')\n    \n#     plt.subplot(2,3,6)\n#     plt.imshow(normal_image); plt.axis('off')\n#     coco.showAnns(annotations)\n#     plt.title('Annotations')\n#     plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# resized = cv2.resize(Masks.data[0].numpy(), (width, height), interpolation = cv2.INTER_AREA)\n# m = results[0].masks.data\n# plt.imshow(Mask[0])\n# print(resized.shape)\n# int(cat_id[4])\n# r.masks.data.shape\n# m.dtype\n# Ids","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def sliding_window(image, kernel_height, y_stride):\n#     y_start, y_end = 0, kernel_height\n#     isheight = 1\n#     while isheight:\n# #         x_start, x_end = 0, kernel_width\n# #         iswidth = 1\n# #         while iswidth:\n#         yield image[y_start:y_end], y_start, y_end\n# #             if  x_end == image.shape[1]:\n# #                 iswidth = 0\n# #             x_start += x_stride\n# #             x_end = min(x_start+kernel_width, image.shape[1])\n#         if y_end == image.shape[0]:\n#             isheight = 0\n#         y_start += y_stride\n#         y_end = min(y_start+kernel_height, image.shape[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import pandas as pd","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# h, w = [], []\n# img = test_coco.getImgIds()\n# for x in img:\n#     d = test_coco.loadImgs(ids=x)\n#     h.append(d[0]['height'])\n#     w.append(d[0]['width'])\n# df = pd.DataFrame({'Id':img, 'Height':h, 'Width':w})\n# df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# df.describe()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# filt = (df['Height'] > 2000) & (df['Width']>2000)\n# df[filt].describe()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# area, cat = [], []\n# annot_id = train_coco.getAnnIds()\n# for x in annot_id:\n#     d = train_coco.loadAnns(ids=x)\n#     cat.append(d[0]['category_id'])\n#     area.append(d[0]['area'])\n# df1 = pd.DataFrame({'Category':cat, 'Area':area})","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# para = df1[(df1['Category'] == 0)]\n# text = df1[(df1['Category'] == 1)]\n# image = df1[(df1['Category'] == 2)]\n# table = df1[(df1['Category'] == 3)]","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}