Translate train.py (and others) to Jupyter notebook

[pedropcamellon]

Hi. I use to work with Jupyter for my ML projects. In order to work with this repo I have to run Python scripts with the corresponding parameters. In the examples provided in Google Colab (https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb) the Python scripts are called like this:

python detect.py --source 0  img.jpg  # image 

I would like to port train.py, and others, to Jupyter notebook extracting what is inside de .py files. I'm currently doing this and want to know if someone is trying the same. Thanks.

[glenn-jocher]

@Pcamellon 👋 Hello! Thanks for asking about handling inference results. YOLOv5 🚀 PyTorch Hub models allow for simple model loading and inference in a pure python environment without using detect.py.

Simple Inference Example

This example loads a pretrained YOLOv5s model from PyTorch Hub as model and passes an image for inference. 'yolov5s' is the lightest and fastest YOLOv5 model. For details on all available models please see the README.

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, custom

# Images
img = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list

# Inference
results = model(img)

# Results
results.print  # or .show(), .save(), .crop(), .pandas(), etc.

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

See YOLOv5 PyTorch Hub Tutorial for details.

Good luck 🍀 and let us know if you have any other questions!