{"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":"import os\nimport cv2\nimport random\nimport numpy as np\nimport pandas as pd\nimport multiprocessing\nfrom tqdm import tqdm\nimport albumentations\n\nimport matplotlib.pyplot as plt\nfrom pylab import rcParams\nrcParams['figure.figsize'] = 24,12","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-15T11:28:34.102348Z","iopub.execute_input":"2022-04-15T11:28:34.102645Z","iopub.status.idle":"2022-04-15T11:28:34.107567Z","shell.execute_reply.started":"2022-04-15T11:28:34.102614Z","shell.execute_reply":"2022-04-15T11:28:34.107004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mnist_dir = '../input/digit-recognizer'\nout_dir = './generated_data'\n\nn_gen = 6  # 6 for debug, 60000 for training yolov5\nos.makedirs(out_dir, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:25:13.323037Z","iopub.execute_input":"2022-04-15T11:25:13.323322Z","iopub.status.idle":"2022-04-15T11:25:13.327711Z","shell.execute_reply.started":"2022-04-15T11:25:13.323291Z","shell.execute_reply":"2022-04-15T11:25:13.326898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train = pd.read_csv(os.path.join(mnist_dir, 'train.csv'))\ndata_train = df_train.values[:, 1:].reshape(-1, 28, 28)\ndata_train = (data_train > 70).astype(int) * 255\nlabel_train = df_train['label'].values","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:25:19.683527Z","iopub.execute_input":"2022-04-15T11:25:19.684369Z","iopub.status.idle":"2022-04-15T11:25:23.611767Z","shell.execute_reply.started":"2022-04-15T11:25:19.684332Z","shell.execute_reply":"2022-04-15T11:25:23.610941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transforms = albumentations.Compose([\n    albumentations.augmentations.geometric.transforms.Affine (scale=None, translate_percent=None, translate_px=0, rotate=0, shear=45, interpolation=1, cval=0, cval_mask=0, mode=3, fit_output=False, always_apply=False, p=1.0),\n    albumentations.ShiftScaleRotate(shift_limit=0.15, scale_limit=(0, 1.), rotate_limit=45, border_mode=4, p=1.0),\n])","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:25:24.742270Z","iopub.execute_input":"2022-04-15T11:25:24.743232Z","iopub.status.idle":"2022-04-15T11:25:24.748006Z","shell.execute_reply.started":"2022-04-15T11:25:24.743169Z","shell.execute_reply":"2022-04-15T11:25:24.747481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generate_one_img(idx):\n    img = np.zeros((4000, 4000)).astype(np.uint8)\n    boxes = []\n    for i1 in range(2):\n        for i2 in range(2):\n            rand_idx = random.randint(0, data_train.shape[0] - 1)\n            rand_digit = data_train[rand_idx]\n            rand_t = 5 if random.random() < 0.3 else 62\n            rand_size = int(random.random() * rand_t) + 1\n            if rand_size < 6 and random.random() < 0.2:\n                rand_digit = cv2.resize(rand_digit.astype(np.uint8).copy(), (16, 16), cv2.INTER_NEAREST)\n            else:\n                rand_digit = rand_digit.repeat(rand_size, axis=0).repeat(rand_size, axis=1)\n\n            label = label_train[rand_idx]\n            l = rand_digit.shape[0]\n\n            x1 = random.randint(2000 * i1, 2000 * i1 + 2000 - l - 1)\n            y1 = random.randint(2000 * i2, 2000 * i2 + 2000 - l - 1)\n            img[y1:y1+l, x1:x1+l] = rand_digit\n            boxes.append([label, x1, x1+l, y1, y1+l])\n\n    bg2 = np.zeros((4000, 4000)).astype(np.uint8)\n    n_ = random.randint(3, 8)\n    for _ in range(n_):\n        mask = np.zeros_like(bg2)\n        if random.random() < 0.5:\n            pt1 = (random.randint(0, 3999), random.randint(0, 3999))\n            pt2 = (random.randint(0, 3999), random.randint(0, 3999))\n            pt3 = (random.randint(0, 3999), random.randint(0, 3999))\n            triangle_cnt = np.array( [pt1, pt2, pt3] )\n            cv2.drawContours(mask, [triangle_cnt], 0, 1, -1)\n        else:\n            pt = (random.randint(0, 3999), random.randint(0, 3999))\n            r = random.randint(128, 1500)\n            cv2.circle(mask, pt, r, 1, -1)\n        bg2[mask > 0.5] = 255 - bg2[mask > 0.5]\n\n    bg = np.zeros((4000, 4000)).astype(np.uint8)\n    bg = (np.random.randint(0, 2, 100).reshape(10,10) * 255).astype(np.uint8)\n    bg = bg.repeat(400, axis=0).repeat(400, axis=1)\n    bg = transforms(image=bg)['image']\n    bg2[bg > 127] = 255 - bg2[bg > 127]\n\n    img[bg2 > 127] = 255 - img[bg2 > 127]\n\n    out_path = os.path.join(out_dir, f'{idx:06d}.jpg')\n    cv2.imwrite(out_path, img)\n    return boxes","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:25:36.050302Z","iopub.execute_input":"2022-04-15T11:25:36.050872Z","iopub.status.idle":"2022-04-15T11:25:36.065172Z","shell.execute_reply.started":"2022-04-15T11:25:36.050825Z","shell.execute_reply":"2022-04-15T11:25:36.064427Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with multiprocessing.Pool(3) as pool:\n    imap = pool.imap(generate_one_img, list(range(n_gen)))\n    boxes = list(tqdm(imap, total=n_gen))","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:26:00.093307Z","iopub.execute_input":"2022-04-15T11:26:00.094161Z","iopub.status.idle":"2022-04-15T11:26:03.466428Z","shell.execute_reply.started":"2022-04-15T11:26:00.094123Z","shell.execute_reply":"2022-04-15T11:26:03.465332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({\n    'id': list(range(n_gen)),\n    'boxes': boxes,\n})","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:26:09.977014Z","iopub.execute_input":"2022-04-15T11:26:09.977370Z","iopub.status.idle":"2022-04-15T11:26:09.986022Z","shell.execute_reply.started":"2022-04-15T11:26:09.977334Z","shell.execute_reply":"2022-04-15T11:26:09.985147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.to_csv('./generated_data.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:26:21.555644Z","iopub.execute_input":"2022-04-15T11:26:21.556426Z","iopub.status.idle":"2022-04-15T11:26:21.565176Z","shell.execute_reply.started":"2022-04-15T11:26:21.556380Z","shell.execute_reply":"2022-04-15T11:26:21.564408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:26:23.937890Z","iopub.execute_input":"2022-04-15T11:26:23.938542Z","iopub.status.idle":"2022-04-15T11:26:23.963504Z","shell.execute_reply.started":"2022-04-15T11:26:23.938503Z","shell.execute_reply":"2022-04-15T11:26:23.962769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize","metadata":{}},{"cell_type":"code","source":"df['filepath'] = df['id'].apply(lambda x: os.path.join(out_dir, f'{x:06d}.jpg'))","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:27:26.786837Z","iopub.execute_input":"2022-04-15T11:27:26.787161Z","iopub.status.idle":"2022-04-15T11:27:26.798452Z","shell.execute_reply.started":"2022-04-15T11:27:26.787123Z","shell.execute_reply":"2022-04-15T11:27:26.797412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(2):\n    f, axarr = plt.subplots(1,3)\n    for p in range(3):\n        idx = i * 3 + p\n        img = cv2.imread(df.loc[idx].filepath)\n        \n        axarr[p].imshow(img)","metadata":{"execution":{"iopub.status.busy":"2022-04-15T11:29:18.627487Z","iopub.execute_input":"2022-04-15T11:29:18.628360Z","iopub.status.idle":"2022-04-15T11:29:35.584833Z","shell.execute_reply.started":"2022-04-15T11:29:18.628310Z","shell.execute_reply":"2022-04-15T11:29:35.583974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}