#clone YOLOv5 and 
!git clone https://github.com/ultralytics/yolov5  # clone repo
%cd yolov5
%pip install -qr requirements.txt # install dependencies
%pip install wandb
import torch
import os
from IPython.display import Image, clear_output  # to display images

print(f"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})")

import wandb

wandb.login()
from google.colab import drive 
drive.mount('/content/drive',force_remount=False)
# set up environment
os.environ["DATASET_DIRECTORY"] = "/content/drive/MyDrive/kickercamds"

#!curl -L "https://app.roboflow.com/ds/x8MCxicjGs?key=RPrzNjhFFB" > roboflow.zip; unzip roboflow.zip -d /content/drive/MyDrive/kickercamds; rm roboflow.zip

"""
Explanation: Custom Training with YOLOv5
In this tutorial, we assemble a dataset and train a custom YOLOv5 model to recognize the objects in our dataset. To do so we will take the following steps:

Gather a dataset of images and label our dataset
Export our dataset to YOLOv5
Train YOLOv5 to recognize the objects in our dataset
Evaluate our YOLOv5 model's performance
Run test inference to view our model at work


Step 1: Install Requirements
End of explanation
"""


!python train.py --noautoanchor --img 428 --rect --batch 128 --epochs 150 --data /content/drive/MyDrive/kickercamds/data.yaml --weights yolov5s.pt --cache

"""
Explanation: Step 3: Train Our Custom YOLOv5 model
Here, we are able to pass a number of arguments:
- img: define input image size
- batch: determine batch size
- epochs: define the number of training epochs. (Note: often, 3000+ are common here!)
- data: Our dataset locaiton is saved in the dataset.location
- weights: specify a path to weights to start transfer learning from. Here we choose the generic COCO pretrained checkpoint.
- cache: cache images for faster training
End of explanation
"""


!python detect.py --weights runs/train/exp2/weights/best.pt --img 428 --conf 0.1 --source /content/drive/MyDrive/kickercamsmall/test/images

#display inference on ALL test images

import glob
from IPython.display import Image, display

for imageName in glob.glob('runs/detect/exp2/*.jpg'): #assuming JPG
    display(Image(filename=imageName))
    print("\n")

"""
Explanation: Run Inference  With Trained Weights
Run inference with a pretrained checkpoint on contents of test/images folder downloaded from Roboflow.
End of explanation
"""


#export your model's weights for future use
from google.colab import files
files.download('./runs/train/exp/weights/best.pt')

"""
Explanation: Conclusion and Next Steps
Congratulations! You've trained a custom YOLOv5 model to recognize your custom objects.
To improve you model's performance, we recommend first interating on your datasets coverage and quality. See this guide for model performance improvement.
To deploy your model to an application, see this guide on exporting your model to deployment destinations.
Once your model is in production, you will want to continually iterate and improve on your dataset and model via active learning.
End of explanation
"""


