{"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":"In this notebook I will show how to use yolov7 with SoccerNet tracking data.  \nAs you will see later, yolov7 is very easy to use, having almost the same interface as yolov5.\n\nNote that it is ambiguous whether SoccerNet datasets can be used in the DFL competition.  \nPlease use these datasets for DFL competition at your own risk.","metadata":{}},{"cell_type":"code","source":"!pip install SoccerNet\n!pip uninstall wandb -y\n\nimport SoccerNet\nfrom SoccerNet.Downloader import SoccerNetDownloader\nimport os\nimport numpy as np\nimport pandas as pd\nfrom IPython.display import Image\nimport cv2\nfrom glob import glob\nimport re\nimport yaml","metadata":{"_kg_hide-input":false,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"yolo_base = '/kaggle/working/yolov7'\nsn_track_base = '/tmp/tracking'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# I will only download the train dataset due to disk space limitations on Kaggle environment.\n# The download script tries to download test and challenge dataset, but it will fail by dummy files.\n# It seems that downloading of the SoccerNet dataset sometimes times out.\n# It may be that bandwidth is limited if we try to download repeatedly.\n!mkdir -p $sn_track_base\n!touch $sn_track_base/test.zip\n!touch $sn_track_base/challenge.zip\n\nmySoccerNetDownloader=SoccerNetDownloader(LocalDirectory=\"/tmp/\")\nmySoccerNetDownloader.downloadDataTask(task=\"tracking\", split=[\"train\", \"test\", \"challenge\"])\n!unzip -q -d $sn_track_base $sn_track_base/train.zip\n!git clone https://github.com/WongKinYiu/yolov7 \n!cd $yolo_base && wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-w6_training.pt","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SoccerNet data setting","metadata":{}},{"cell_type":"code","source":"# Of the 57 half , 50 are used for training and 7 for validation.\ntrn_dirs = sorted(glob(f'{sn_track_base}/train/SNMOT*'))\n#tst_dirs = sorted(glob(f'{sn_track_base}/test/SNMOT*'))\ntst_dirs = trn_dirs[50:]\ntrn_dirs = trn_dirs[:50]\n\ndebug=True\nif debug:\n    trn_dirs = trn_dirs[:5]\n    tst_dirs = tst_dirs[:5]\n\nsplit_dirs = {\n    'train':trn_dirs,\n    'valid':tst_dirs\n}\n\n# ４ types of labels are used.\nlabels = ['ball', 'player', 'referee', 'goalkeepers']\nlabel_dict = {'ball': 0, 'player': 1, 'referee': 2, 'goalkeeper': 3, 'goalkeepers': 3}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# yolov7 setting\nOnly three types of files need to be prepared:  \n(1) image files  \n(2) label files  \n(3) configuration files","metadata":{}},{"cell_type":"code","source":"#(1) image file path\nyolo_train_img_dir = f'{yolo_base}/images/train'\nyolo_valid_img_dir = f'{yolo_base}/images/valid'\n\n#(2) label file path\nyolo_train_label_dir = f'{yolo_base}/labels/train'\nyolo_valid_label_dir = f'{yolo_base}/labels/valid'\n\n#(3) config file path\nyaml_file = f'{yolo_base}/data.yaml'\n\n#!rm -rf /home/tito/kaggle/dfl-bundesliga-data-shootout/work/yolov5\nos.makedirs(yolo_train_img_dir, exist_ok=True)\nos.makedirs(yolo_valid_img_dir, exist_ok=True)\nos.makedirs(yolo_train_label_dir, exist_ok=True)\nos.makedirs(yolo_valid_label_dir, exist_ok=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# convert from x,y,w,h to yolo format\ndef get_yolo_format_bbox(img_w, img_h, box):\n    w = box[2]\n    h = box[3]\n    xc = box[0] + int(np.round(w/2))\n    yc = box[1] + int(np.round(h/2))\n    box = [xc/img_w, yc/img_h, w/img_w, h/img_h]\n    box = [f\"{i:.4g}\" for i in box]\n    return box\n    \n# get SoccerNet label info \ndef get_info(info):\n    results = []\n    for line in open(info):\n        m = re.match('trackletID_(\\d+)= (\\S*).*', line.replace(';', ' '))\n        if m:\n            if m.group(2) not in label_dict:\n                #print('bad label:', m.group(2))\n                continue \n            results.append([m.group(1), m.group(2)])\n    return pd.DataFrame(results, columns=['id','lbl']).set_index('id').to_dict()['lbl']\n\n# make image simlink and label files\nfor split in split_dirs:\n    if split == 'train':\n        yolo_img_dir = yolo_train_img_dir\n        yolo_label_dir = yolo_train_label_dir\n    else:\n        yolo_img_dir = yolo_valid_img_dir\n        yolo_label_dir = yolo_valid_label_dir\n        \n    for this_dir in split_dirs[split]:\n        print('this_dir:',this_dir)\n        video = this_dir.split('/')[-1]\n        info = this_dir + '/gameinfo.ini'\n        det = this_dir + '/gt/gt.txt'\n        info_dict = get_info(info)\n        det_df = pd.read_csv(det, names=['frame','player','x','y','w','h','f1','f2','f3','f4'], usecols=['frame','player','x','y','w','h'])\n        det_df['label'] = det_df.player.astype(str).map(info_dict)\n        det_df['label_id'] = det_df['label'].map(label_dict)\n        # check\n        ng_list = list(det_df[det_df.label_id.isnull()].label.unique())\n        if len(ng_list)>0:\n            #print('ng_list:',ng_list, det_df.dropna().shape, det_df.shape)\n            det_df = det_df.dropna()\n        for grp, grp_df in det_df.groupby('frame'):\n            frame = f'{grp:06}'\n            img_file = f'{this_dir}/img1/{frame}.jpg'\n            dst_file = f'{yolo_img_dir}/{video}_{frame}.jpg'\n            if not os.path.exists(dst_file):\n                os.symlink(img_file, dst_file)\n                #print(img_file)\n            img = cv2.imread(dst_file)\n            height, width, _ = img.shape \n            bboxes = []\n            for arr in grp_df[['x', 'y', 'w', 'h', 'label_id']].values:\n                box = get_yolo_format_bbox(width, height, arr[:4])\n                bboxes.append([arr[4]]+box)\n            file_name = f'{yolo_label_dir}/{video}_{frame}.txt'\n            with open(file_name, 'w') as f:\n                for i, bbox in enumerate(bboxes):\n                    bbox = [str(i) for i in bbox]\n                    str_bbox = ' '.join(bbox)\n                    f.write(str_bbox)\n                    f.write('\\n')","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dump config file\ndata_yaml = dict(\n    train = yolo_train_img_dir,\n    val = yolo_valid_img_dir,\n    nc = 4,\n    names = labels\n)\n\nwith open(yaml_file, 'w') as outfile:\n    yaml.dump(data_yaml, outfile, default_flow_style=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train!\n!cd $yolo_base && time python train_aux.py \\\n    --workers 2 \\\n    --batch-size 4 \\\n    --data data.yaml \\\n    --img 1280 \\\n    --epochs 10 \\\n    --cfg cfg/training/yolov7-w6.yaml \\\n    --weights yolov7-w6_training.pt \\\n    --name yolov7-w6-custom \\\n    --hyp data/hyp.scratch.custom.yaml","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# predict!\ncheck_images = f'{tst_dirs[0]}/img1'\nweight_file = 'runs/train/yolov7-w6-custom/weights/best.pt'\n!cd $yolo_base && python detect.py \\\n    --weights $weight_file \\\n    --source  $check_images \\\n    --img-size 1280 \\\n    --project dfl","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check some results. \ndisplay(Image(f'{yolo_base}/dfl/exp/000200.jpg'))\ndisplay(Image(f'{yolo_base}/dfl/exp/000400.jpg'))\ndisplay(Image(f'{yolo_base}/dfl/exp/000600.jpg'))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}