notebooks/ready-to-run-rico-ui-detection.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install ultralytics","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-28T22:00:33.556327Z","iopub.execute_input":"2023-03-28T22:00:33.556777Z","iopub.status.idle":"2023-03-28T22:00:52.812879Z","shell.execute_reply.started":"2023-03-28T22:00:33.556735Z","shell.execute_reply":"2023-03-28T22:00:52.811194Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stdout","text":"Collecting ultralytics\n  Downloading ultralytics-8.0.58-py3-none-any.whl (486 kB)\n\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m486.8/486.8 kB\u001b[0m \u001b[31m11.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n\u001b[?25hCollecting thop>=0.1.1\n  Downloading thop-0.1.1.post2209072238-py3-none-any.whl (15 kB)\nRequirement already satisfied: torchvision>=0.8.1 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (0.14.0+cpu)\nRequirement already satisfied: Pillow>=7.1.2 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (9.4.0)\nRequirement already satisfied: torch>=1.7.0 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (1.13.0+cpu)\nRequirement already satisfied: tqdm>=4.64.0 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (4.64.1)\nRequirement already satisfied: sentry-sdk in /opt/conda/lib/python3.7/site-packages (from ultralytics) (1.15.0)\nRequirement already satisfied: PyYAML>=5.3.1 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (6.0)\nRequirement already satisfied: psutil in /opt/conda/lib/python3.7/site-packages (from ultralytics) (5.9.3)\nRequirement already satisfied: scipy>=1.4.1 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (1.7.3)\nRequirement already satisfied: matplotlib>=3.2.2 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (3.5.3)\nCollecting opencv-python>=4.6.0\n  Downloading opencv_python-4.7.0.72-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (61.8 MB)\n\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m61.8/61.8 MB\u001b[0m \u001b[31m17.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n\u001b[?25hRequirement already satisfied: pandas>=1.1.4 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (1.3.5)\nRequirement already satisfied: requests>=2.23.0 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (2.28.2)\nRequirement already satisfied: seaborn>=0.11.0 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (0.12.2)\nRequirement already satisfied: numpy>=1.21.6 in /opt/conda/lib/python3.7/site-packages (from ultralytics) (1.21.6)\nRequirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (1.4.4)\nRequirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (2.8.2)\nRequirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (0.11.0)\nRequirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (4.38.0)\nRequirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (23.0)\nRequirement already satisfied: pyparsing>=2.2.1 in /opt/conda/lib/python3.7/site-packages (from matplotlib>=3.2.2->ultralytics) (3.0.9)\nRequirement already satisfied: pytz>=2017.3 in /opt/conda/lib/python3.7/site-packages (from pandas>=1.1.4->ultralytics) (2022.7.1)\nRequirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests>=2.23.0->ultralytics) (2022.12.7)\nRequirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests>=2.23.0->ultralytics) (1.26.14)\nRequirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests>=2.23.0->ultralytics) (3.4)\nRequirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.7/site-packages (from requests>=2.23.0->ultralytics) (2.1.1)\nRequirement already satisfied: typing_extensions in /opt/conda/lib/python3.7/site-packages (from seaborn>=0.11.0->ultralytics) (4.4.0)\nRequirement already satisfied: six>=1.5 in /opt/conda/lib/python3.7/site-packages (from python-dateutil>=2.7->matplotlib>=3.2.2->ultralytics) (1.16.0)\nInstalling collected packages: opencv-python, thop, ultralytics\n  Attempting uninstall: opencv-python\n    Found existing installation: opencv-python 4.5.4.60\n    Uninstalling opencv-python-4.5.4.60:\n      Successfully uninstalled opencv-python-4.5.4.60\nSuccessfully installed opencv-python-4.7.0.72 thop-0.1.1.post2209072238 ultralytics-8.0.58\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m","output_type":"stream"}]},{"cell_type":"code","source":"!pip freeze > requirements.txt  ","metadata":{"execution":{"iopub.status.busy":"2023-03-28T22:01:31.187263Z","iopub.execute_input":"2023-03-28T22:01:31.187827Z","iopub.status.idle":"2023-03-28T22:01:34.121426Z","shell.execute_reply.started":"2023-03-28T22:01:31.187774Z","shell.execute_reply":"2023-03-28T22:01:34.119762Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"code","source":"import os\nimport shutil\nimport random\nimport json\nimport glob\nimport torch\n\nfrom ultralytics import YOLO\nfrom PIL import Image","metadata":{"execution":{"iopub.status.busy":"2023-03-28T22:00:52.815452Z","iopub.execute_input":"2023-03-28T22:00:52.816405Z","iopub.status.idle":"2023-03-28T22:00:53.541868Z","shell.execute_reply.started":"2023-03-28T22:00:52.816366Z","shell.execute_reply":"2023-03-28T22:00:53.540406Z"},"trusted":true},"execution_count":2,"outputs":[]},{"cell_type":"code","source":"print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_folder = \"/kaggle/input/ui-kits-simple/rico_dataset_v0.1_semantic_annotations/semantic_annotations\"\noutput_image_folder = \"/kaggle/working/data/rico/images\"\noutput_json_folder = \"/kaggle/working/data/rico/jsons\"\n\nos.makedirs(output_image_folder, exist_ok=True)\nos.makedirs(output_json_folder, exist_ok=True)\n\nfor file_name in os.listdir(input_folder):\n    if file_name.endswith(\".png\") or file_name.endswith(\".jpg\"):\n        shutil.copy(os.path.join(input_folder, file_name), os.path.join(output_image_folder, file_name))\n    elif file_name.endswith(\".json\"):\n        shutil.copy(os.path.join(input_folder, file_name), os.path.join(output_json_folder, file_name))\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#bounds => [left, top, right, bottom]\n\ninput_path = \"/kaggle/working/data/rico/jsons\"\noutput_path = '/kaggle/working/data/rico/txt'\n\nfor filename in os.listdir(input_path):\n    with open(os.path.join(input_path, filename)) as f:\n        data = json.load(f)\n\n        def get_child_components(data, labels=[]):\n            if 'children' in list(data.keys()):\n                for c in data['children']:\n\n                    label_map = {\n                        'Image': 0,\n                        'Text': 1,\n                        'Text Button': 2,\n                        'Input': 3,\n                        'Icon': 4,\n                        'List Item': 5,\n                        'Another': 6,\n                        # more mappings if needed\n                    }\n\n                    #Convert these bounds to YOLO format\n\n                    width = c['bounds'][2] - c['bounds'][0]\n                    height = c['bounds'][3] - c['bounds'][1]\n\n                    x_center = c['bounds'][0] + width/2\n                    y_center = c['bounds'][1] + height/2\n                    \n                    base = data['bounds']\n                    \n                    image_width = 1440 #int(base[2])\n                    image_height = 2560 #int(base[3])\n                    \n                    x_center_norm = x_center/image_width\n                    y_center_norm = y_center/image_height\n                    width_norm = width/image_width\n                    height_norm = height/image_height\n\n                    component_index = label_map.get(c['componentLabel'], 0)\n\n                    yolo_format = [component_index, x_center_norm, y_center_norm, width_norm, height_norm]\n\n                    labels.append({c['componentLabel']: yolo_format})\n\n                    get_child_components(c, labels)\n            return labels\n\n        try:\n            os.mkdir(output_path)\n        except:\n            pass\n        \n        output_file_path = os.path.join(output_path, filename[:-5] + '.txt')\n        \n        with open(output_file_path, 'w') as f:\n            for i in get_child_components(data):\n                for k in i.keys():\n                    f.write('{} {} {} {} {}\\n'.format(int(i[k][0]), float(i[k][1]), float(i[k][2]), float(i[k][3]), float(i[k][4])))\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"Finds corresponding png files after distributed txt files\"\"\"\n\n# Define the paths to the image and text files\nimage_path = \"/kaggle/working/data/rico/images\"\ntext_path = \"/kaggle/working/data/rico/txt\"\n\n# Define the output directories for the train, val, and test sets\ntrain_path = \"/kaggle/working/data/rico/train\"\nval_path = \"/kaggle/working/data/rico/val\"\ntest_path = \"/kaggle/working/data/rico/test\"\n\n# Define the ratio of the train and val sets\ntrain_ratio = 0.70\nval_ratio = 0.15\ntest_ratio = 0.15\n\n# Get the list of text files\ntext_files = sorted(os.listdir(text_path))\n\n# Shuffle the list of text files\nrandom.shuffle(text_files)\n\n# Calculate the number of files for each set\nnum_train = int(len(text_files) * train_ratio)\nnum_val = int(len(text_files) * val_ratio)\nnum_test = int(len(text_files) - num_train - num_val)\n\n# Assign each file to a set\ntrain_files = text_files[:num_train]\nval_files = text_files[num_train:num_train+num_val]\ntest_files = text_files[num_train+num_val:]\n\n# Create the output directories, if exists delete them and create new ones\nif os.path.exists(train_path):\n    shutil.rmtree(train_path)\nos.makedirs(train_path)\n\nif os.path.exists(val_path):\n    shutil.rmtree(val_path)\nos.makedirs(val_path)\n\nif os.path.exists(test_path):\n    shutil.rmtree(test_path)\nos.makedirs(test_path)\n\n# Copy train txt files to the train directory, then copy the corresponding png files\nfor file in train_files:\n    shutil.copy(os.path.join(text_path, file), train_path)\n    shutil.copy(os.path.join(image_path, file[:-4] + \".png\"), train_path)\n\n# Copy val txt files to the val directory, then copy the corresponding png files\nfor file in val_files:\n    shutil.copy(os.path.join(text_path, file), val_path)\n    shutil.copy(os.path.join(image_path, file[:-4] + \".png\"), val_path)\n\n# Copy test txt files to the test directory, then copy the corresponding png files\nfor file in test_files:\n    shutil.copy(os.path.join(text_path, file), test_path)\n    shutil.copy(os.path.join(image_path, file[:-4] + \".png\"), test_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Print the number of files in each set, both text and image files\nprint(\"Number of train files: \", len(os.listdir(train_path)))\nprint(\"Number of val files: \", len(os.listdir(val_path)))\nprint(\"Number of test files: \", len(os.listdir(test_path)))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load a model\n# model_path = \"/kaggle/input/core-components/model.yaml\"\n\nmodel = YOLO('yolov8n.yaml') # build a new model from YAML\nmodel = YOLO('yolov8n.pt',task='detect')  # load a pretrained model (recommended for training)\n# model = YOLO(yaml_path, task='detect').load('yolov8n.pt') # build from YAML and transfer weights","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train the model\nyaml_path = \"/kaggle/input/core-components/data.yaml\"\ndevice = 'cuda:0'  # Use the first available GPU\nmodel.train(data=yaml_path, epochs=2, imgsz=640, device=device, val=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Export the model as torchscript\nbest_weight_path = '/kaggle/working/runs/detect/train/weights/best.pt'\nmodel = YOLO(best_weight_path) \nmodel.fuse()  \nmodel.info(verbose=True)  # Print model information\nmodel.export()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test model with a random screenshot\ndir_path = \"/kaggle/input/ui-kits-simple/unique_uis/combined\"\njpg_files = glob.glob(dir_path + \"/*.jpg\")\n\nif jpg_files:\n    try:\n        random_jpg = random.choice(jpg_files)\n        image = Image.open(random_jpg)\n        result = model.predict(source=image, save=True)\n        print(result)\n        \n    except Exception as e:\n        print(\"Error occurred while processing image:\", e)\n        \nelse:\n    print(\"No JPG files found in directory.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import shutil, os\n\n# def kaggle_export(output_dir, zip_name):\n#     \"\"\"Create a zip file and export your work in the kaggle working directory.\n\n#     Args:\n#         output_dir (str): The path to the directory to be zipped.\n#         zip_name (str): The name of the zip file to be created.\n\n#     Returns:\n#         None\n#     \"\"\"\n#     if not os.path.exists(output_dir):\n#         print(f\"[Error]: output directory '{output_dir}' does not exist\")\n    \n#     else:\n#         try:\n#             shutil.make_archive(zip_name, 'zip', output_dir)\n#             print(f\"[Success]: Your zip file is successfully created!\")\n\n#         except Exception as e:\n#             print(f\"[Error]: An error occurred while creating zip: {e}\")\n\n# kaggle_export(\"/kaggle/working/runs/detect/train\", \"model_result\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}