36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
from ultralytics import YOLO
|
|
from pathlib import Path
|
|
|
|
|
|
def foo():
|
|
model = YOLO("yolo11x.pt")
|
|
example = Path(
|
|
'/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
|
|
results = model.predict(example)
|
|
print('print done predicting')
|
|
print(f'done predicting; {len(results)} results')
|
|
for res in results:
|
|
print(f" >>> {res.boxes.data}")
|
|
|
|
def main():
|
|
# Load a model
|
|
model = YOLO("yolo11s.pt") # load a pretrained model (recommended for training)
|
|
|
|
# Train the model
|
|
train = '/home/didier/all-dressed-programming/dodgeball-tracker/data/wheels_004.yaml'
|
|
model.train(data=Path(train),
|
|
epochs=150,
|
|
imgsz=640,
|
|
cache='ram')
|
|
print('Done training')
|
|
bad_example = Path("/home/didier/all-dressed-programming/dodgeball-tracker/data/label-studio-exports/2024-11-12_001/images/train/2a6267f7-PB120166.JPG")
|
|
example = Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
|
|
results = model.predict(bad_example)
|
|
print(f'done predicting; {len(results)} results')
|
|
for res in results:
|
|
print(f" >>> {res.boxes.data}")
|
|
|
|
if __name__ == '__main__':
|
|
foo()
|
|
|