{"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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":40501,"databundleVersionId":5285958,"sourceType":"competition"}],"dockerImageVersionId":30396,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"pip install ultralytics==8.0.55","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2025-10-12T11:49:42.867461Z","iopub.execute_input":"2025-10-12T11:49:42.868622Z","iopub.status.idle":"2025-10-12T11:49:51.223013Z","shell.execute_reply.started":"2025-10-12T11:49:42.868583Z","shell.execute_reply":"2025-10-12T11:49:51.221614Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom sklearn.model_selection import train_test_split\nfrom ultralytics import YOLO","metadata":{"execution":{"iopub.status.busy":"2025-10-12T11:50:53.852512Z","iopub.execute_input":"2025-10-12T11:50:53.852885Z","iopub.status.idle":"2025-10-12T11:50:53.857628Z","shell.execute_reply.started":"2025-10-12T11:50:53.852842Z","shell.execute_reply":"2025-10-12T11:50:53.856608Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-10-12T11:36:24.793544Z","iopub.execute_input":"2025-10-12T11:36:24.794385Z","iopub.status.idle":"2025-10-12T11:36:24.819322Z","shell.execute_reply.started":"2025-10-12T11:36:24.794347Z","shell.execute_reply":"2025-10-12T11:36:24.818421Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-10-12T11:36:31.556479Z","iopub.execute_input":"2025-10-12T11:36:31.556849Z","iopub.status.idle":"2025-10-12T11:36:31.565590Z","shell.execute_reply.started":"2025-10-12T11:36:31.556818Z","shell.execute_reply":"2025-10-12T11:36:31.564889Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-10-12T11:36:34.089487Z","iopub.execute_input":"2025-10-12T11:36:34.090112Z","iopub.status.idle":"2025-10-12T11:36:34.095521Z","shell.execute_reply.started":"2025-10-12T11:36:34.090075Z","shell.execute_reply":"2025-10-12T11:36:34.094528Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load a pretrained model","metadata":{}},{"cell_type":"code","source":"model = YOLO(\"yolov8n.pt\")","metadata":{"execution":{"iopub.status.busy":"2025-10-12T11:36:36.691278Z","iopub.execute_input":"2025-10-12T11:36:36.692033Z","iopub.status.idle":"2025-10-12T11:36:37.289064Z","shell.execute_reply.started":"2025-10-12T11:36:36.692001Z","shell.execute_reply":"2025-10-12T11:36:37.288293Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"model.train(\n    data=yaml_path,\n    epochs=100,\n    device=\"cuda:0\",\n    project=\"dtp-f\")","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2025-10-12T11:36:53.343606Z","iopub.execute_input":"2025-10-12T11:36:53.344588Z","iopub.status.idle":"2025-10-12T11:45:29.342710Z","shell.execute_reply.started":"2025-10-12T11:36:53.344551Z","shell.execute_reply":"2025-10-12T11:45:29.341882Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":"2025-10-12T11:46:02.562634Z","iopub.execute_input":"2025-10-12T11:46:02.563096Z","iopub.status.idle":"2025-10-12T11:46:53.662438Z","shell.execute_reply.started":"2025-10-12T11:46:02.563056Z","shell.execute_reply":"2025-10-12T11:46:53.661391Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":{}}]}