{"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":"# Imports","metadata":{}},{"cell_type":"code","source":"!pip install ultralytics","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-03-22T10:58:30.923915Z","iopub.execute_input":"2023-03-22T10:58:30.924334Z","iopub.status.idle":"2023-03-22T10:58:46.230687Z","shell.execute_reply.started":"2023-03-22T10:58:30.924249Z","shell.execute_reply":"2023-03-22T10:58:46.229518Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom sklearn.model_selection import train_test_split\nfrom ultralytics import YOLO","metadata":{"execution":{"iopub.status.busy":"2023-03-22T10:58:46.232973Z","iopub.execute_input":"2023-03-22T10:58:46.233358Z","iopub.status.idle":"2023-03-22T10:58:48.819745Z","shell.execute_reply.started":"2023-03-22T10:58:46.233323Z","shell.execute_reply":"2023-03-22T10:58:48.818702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset set up","metadata":{}},{"cell_type":"markdown","source":"## Make a validation set","metadata":{}},{"cell_type":"code","source":"initial_train = os.listdir(\"/kaggle/input/ada-image-recognition-fiber/dataset/images/train\")\ntrain_imgs, val_imgs = train_test_split(initial_train, train_size=0.9, random_state=1)","metadata":{"execution":{"iopub.status.busy":"2023-03-22T10:58:48.821199Z","iopub.execute_input":"2023-03-22T10:58:48.821912Z","iopub.status.idle":"2023-03-22T10:58:48.978426Z","shell.execute_reply.started":"2023-03-22T10:58:48.821839Z","shell.execute_reply":"2023-03-22T10:58:48.977409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_txt = \"train.txt\"\nval_txt = \"val.txt\"\ntrain_dir = \"/kaggle/input/ada-image-recognition-fiber/dataset/images/train\"\nfor imgs, txt in zip([train_imgs, val_imgs], [train_txt, val_txt]):\n    with open(txt, 'w') as file:\n        for img in imgs:\n            file.write(os.path.join(train_dir, img)+\"\\n\")","metadata":{"execution":{"iopub.status.busy":"2023-03-22T10:58:48.981149Z","iopub.execute_input":"2023-03-22T10:58:48.981517Z","iopub.status.idle":"2023-03-22T10:58:48.992000Z","shell.execute_reply.started":"2023-03-22T10:58:48.981472Z","shell.execute_reply":"2023-03-22T10:58:48.991085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create a custom .yaml file for training","metadata":{}},{"cell_type":"markdown","source":"Create a new `custom.yaml` file containing the dataset information required by YOLO models.","metadata":{}},{"cell_type":"code","source":"yaml_path = \"/kaggle/working/custom.yaml\"\nyaml_content = f\"\"\" \ntrain: {train_txt}\nval: {val_txt}\ntest: /kaggle/input/ada-image-recognition-fiber/dataset/images/test\n\n# class names\nnames: \n  0: Screw,\n  1: Foam\n  2: Plastic cover\n  3: Tie-wrap, \n  4: Rubbers\n\"\"\"\n\nwith open(yaml_path, 'w') as file:\n    file.write(yaml_content)","metadata":{"execution":{"iopub.status.busy":"2023-03-22T10:58:48.993446Z","iopub.execute_input":"2023-03-22T10:58:48.994064Z","iopub.status.idle":"2023-03-22T10:58:48.999673Z","shell.execute_reply.started":"2023-03-22T10:58:48.994029Z","shell.execute_reply":"2023-03-22T10:58:48.998611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load a pretrained model","metadata":{}},{"cell_type":"code","source":"model = YOLO(\"yolov8n.pt\")","metadata":{"execution":{"iopub.status.busy":"2023-03-22T10:58:49.001119Z","iopub.execute_input":"2023-03-22T10:58:49.001511Z","iopub.status.idle":"2023-03-22T10:58:49.373135Z","shell.execute_reply.started":"2023-03-22T10:58:49.001463Z","shell.execute_reply":"2023-03-22T10:58:49.372172Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"model.train(\n    data=yaml_path,\n    epochs=3,\n    device=\"cuda:0\",\n    project=\"dtp-f\")","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-03-22T10:58:49.374775Z","iopub.execute_input":"2023-03-22T10:58:49.375148Z","iopub.status.idle":"2023-03-22T11:23:58.650599Z","shell.execute_reply.started":"2023-03-22T10:58:49.375112Z","shell.execute_reply":"2023-03-22T11:23:58.649422Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Information about the training can be found in *runs/detect/runs/detect/train*. The next training will get saved to *runs/detect/runs/detect/train2/* and so on.","metadata":{}},{"cell_type":"markdown","source":"# Predict","metadata":{}},{"cell_type":"code","source":"trained_model = YOLO(\"/kaggle/working/dtp-f/train/weights/best.pt\")\nresults = trained_model.predict(\n    source=\"/kaggle/input/ada-image-recognition-fiber/dataset/images/test\",\n    device=\"0\",\n    save=True, \n    save_txt=True,\n    project=\"dtp-f\")  # save predictions as labels","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-03-22T11:23:58.652148Z","iopub.execute_input":"2023-03-22T11:23:58.652829Z","iopub.status.idle":"2023-03-22T11:25:55.611017Z","shell.execute_reply.started":"2023-03-22T11:23:58.652789Z","shell.execute_reply":"2023-03-22T11:25:55.609862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"As you can read in the above cell, predictions have been written to disk. You can find the predicted labels at *dtp-f/predict/labels*. You can now use that to create a submission file in the expected format.","metadata":{}},{"cell_type":"markdown","source":"# Next steps\nAs you have seen, it doesn't take many lines of code to train an object detection model for a custom dataset. But a lot of work still remains to reach the desired results. Here are a few pointers on what to try to improve results.\n\n- Investigate training images.\n- Try out image augmentation\n- Is your validation set well representative of the test set ?\n- Try other models","metadata":{}}]}