combat/operators/preproc.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "raw_datasets = Path(\"datasets/battlefield\")\n",
    "output_dir = Path(\"datasets\")\n",
    "\n",
    "def create_label_file():\n",
    "    label = f\"\"\"path: .\n",
    "train: train_all_in_one.txt\n",
    "val: val_all_in_one.txt\n",
    "# test: test_all_in_one.txt\n",
    "\n",
    "names:\n",
    "  0: operator\n",
    "\"\"\"\n",
    "    with open(output_dir / \"label.yaml\", \"w\") as f:\n",
    "        f.write(label)\n",
    "\n",
    "create_label_file()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def move_draw():\n",
    "    draw_list = raw_datasets.glob(\"*draw.png\")\n",
    "    draw_dir = raw_datasets / \"draw\"\n",
    "    draw_dir.mkdir(parents=True, exist_ok=True)\n",
    "    for draw in draw_list:\n",
    "        draw.rename(draw_dir / draw.name)\n",
    "\n",
    "move_draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000 / 12903\n",
      "2000 / 12903\n",
      "3000 / 12903\n",
      "4000 / 12903\n",
      "5000 / 12903\n",
      "6000 / 12903\n",
      "7000 / 12903\n",
      "8000 / 12903\n",
      "9000 / 12903\n",
      "10000 / 12903\n",
      "11000 / 12903\n",
      "12000 / 12903\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "import random\n",
    "import pypinyin\n",
    "\n",
    "zoom = 640.0 / 720.0\n",
    "\n",
    "def to_pinyin(text):\n",
    "    return ''.join([item[0] for item in pypinyin.pinyin(\n",
    "            text, style=pypinyin.NORMAL, errors='default')]).lower()\n",
    "\n",
    "def pre_proc():\n",
    "    proc_images = raw_datasets.parent / \"images\"\n",
    "    proc_images.mkdir(parents=True, exist_ok=True)\n",
    "    proc_label = raw_datasets.parent / \"labels\"\n",
    "    proc_label.mkdir(parents=True, exist_ok=True)\n",
    "    right_move = 1138 - 640\n",
    "\n",
    "    data_list = list(raw_datasets.glob(\"*.png\"))\n",
    "    for idx, image_path in enumerate(data_list):\n",
    "        label_file = image_path.with_suffix(\".txt\")\n",
    "        lines = None\n",
    "        if label_file.exists():\n",
    "            with open(label_file, \"r\") as f:\n",
    "                lines = f.readlines()\n",
    "            if len(lines) < 2:\n",
    "                continue\n",
    "\n",
    "        if lines:\n",
    "            stage_name = to_pinyin(lines[0].strip())\n",
    "        else:\n",
    "            stage_name = \"BG\"\n",
    "        \n",
    "        left_img_path = proc_images / f\"{stage_name}_{image_path.name}_left.png\"\n",
    "        right_img_path = proc_images / f\"{stage_name}_{image_path.name}_right.png\"\n",
    "        if not left_img_path.exists() or not right_img_path.exists():\n",
    "            image = Image.open(image_path)\n",
    "            image = image.convert(\"RGB\")\n",
    "            image = image.resize((1138, 640))\n",
    "\n",
    "            if not left_img_path.exists():\n",
    "                left_box = (0, 0, 640, 640)\n",
    "                left_image = image.crop(left_box)\n",
    "                left_image.save(left_img_path)\n",
    "\n",
    "            if not right_img_path.exists():\n",
    "                right_box = (right_move, 0, 1138, 640)\n",
    "                right_image = image.crop(right_box)\n",
    "                right_image.save(right_img_path)\n",
    "\n",
    "\n",
    "        left_label = \"\"\n",
    "        right_label = \"\"\n",
    "        if lines:\n",
    "            for line_idx in range(1, len(lines)):\n",
    "                oper_line = lines[line_idx]\n",
    "                name, x, y, skill = oper_line.strip().split(\" \")\n",
    "                skill_y = (float(y) - 108) * zoom\n",
    "                x = float(x) * zoom\n",
    "                y = float(y) * zoom\n",
    "                left = x - random.randint(42, 46)\n",
    "                right = x + random.randint(42, 46)\n",
    "                x = (left + right) / 2.0\n",
    "                top = y + random.randint(20, 25)\n",
    "                bottom = y + random.randint(50, 55)\n",
    "                y = (top + bottom) / 2.0\n",
    "\n",
    "                skill: bool = bool(int(skill))\n",
    "\n",
    "                if x < 640:\n",
    "                    x = x / 640.0\n",
    "                    y = y / 640.0\n",
    "                    w = (right - left) / 640.0\n",
    "                    h = (bottom - top) / 640.0\n",
    "                    left_label += f\"0 {x} {y} {w} {h}\\n\"\n",
    "                    if skill:\n",
    "                        skill_y = skill_y / 640.0\n",
    "                        # left_label += f\"1 {x} {skill_y} 0.0875 0.0875\\n\"\n",
    "                else:\n",
    "                    x = (x - right_move) / 640.0\n",
    "                    y = y / 640.0\n",
    "                    w = (right - left) / 640.0\n",
    "                    h = (bottom - top) / 640.0\n",
    "                    right_label += f\"0 {x} {y} {w} {h}\\n\"\n",
    "                    if skill:\n",
    "                        skill_y = skill_y / 640.0\n",
    "                        # right_label += f\"1 {x} {skill_y} 0.0875 0.0875\\n\"\n",
    "\n",
    "        with open(proc_label / f\"{stage_name}_{image_path.name}_left.txt\", \"w\") as f:\n",
    "            f.write(left_label)\n",
    "        with open(proc_label / f\"{stage_name}_{image_path.name}_right.txt\", \"w\") as f:\n",
    "            f.write(right_label)\n",
    "\n",
    "        if idx % 1000 == 0:\n",
    "            print(f\"{idx} / {len(data_list)}\")\n",
    "\n",
    "pre_proc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_train_and_val():\n",
    "    proc_images = raw_datasets.parent / \"images\"\n",
    "    image_list = list(proc_images.glob(\"*.png\"))\n",
    "    \n",
    "    train_sets = []\n",
    "    val_sets = []\n",
    "    for idx, image in enumerate(image_list):\n",
    "        if idx % 100 < 20:\n",
    "            val_sets.append(image)\n",
    "        else:\n",
    "            train_sets.append(image)\n",
    "    \n",
    "    train_label = \"\"\n",
    "    val_label = \"\"\n",
    "    for image in train_sets:\n",
    "        train_label += f\"{image.absolute()}\\n\"\n",
    "    for image in val_sets:\n",
    "        val_label += f\"{image.absolute()}\\n\"\n",
    "    \n",
    "    with open(output_dir / \"train_all_in_one.txt\", \"w\") as f:\n",
    "        f.write(train_label)\n",
    "    with open(output_dir / \"val_all_in_one.txt\", \"w\") as f:\n",
    "        f.write(val_label)\n",
    "\n",
    "split_train_and_val()\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yolov5",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}