Mercurial > repos > bgruening > json2yolosegment

annotate yolov8.py @ 3:97bc82ee2a61 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 743c8acf1ea4e4b1e718743d3772b7e592646611
author bgruening
date Mon, 14 Jul 2025 18:28:46 +0000
parents 158e6ce48345
children f6990d85161c
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
1 import argparse
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
2 import os
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
3 import pathlib
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
4 import time
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
5 from argparse import RawTextHelpFormatter
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
6 from collections import defaultdict
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
7
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
8 import cv2
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
9 import numpy as np
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
10 from termcolor import colored
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
11 from tifffile import imwrite
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
12 from ultralytics import YOLO
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
13
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
14
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
15 #
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
16 # Input arguments
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
17 #
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
18 parser = argparse.ArgumentParser(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
19 description='train/predict dataset with YOLOv8',
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
20 epilog="""USAGE EXAMPLE:\n\n~~~~Prediction~~~~\n\
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
21 python yolov8.py --test_path=/g/group/user/data --model_path=/g/cba/models --model_name=yolov8n --save_dir=/g/group/user/results --iou=0.7 --confidence=0.5 --image_size=320 --run_dir=/g/group/user/runs --foldername=batch --headless --num_classes=1 max_det=1 --class_names_file=/g/group/user/class_names.txt\n\
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
22 \n~~~~Training~~~~ \n\
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
23 python yolov8.py --train --yaml_path=/g/group/user/example.yaml --model_path=/g/cba/models --model_name=yolov8n --run_dir=/g/group/user/runs/ --image_size=320 --epochs=150 --scale=0.3 --hsv_v=0.5 --model_format=pt --degrees=180 --class_names_file=/g/group/user/class_names.txt""", formatter_class=RawTextHelpFormatter)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
24 parser.add_argument("--dir_path",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
25 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
26 "Path to the training data directory."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
27 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
28 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
29 parser.add_argument("--yaml_path",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
30 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
31 "YAML file with all the data paths"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
32 " i.e. for train, test, valid data."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
33 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
34 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
35 parser.add_argument("--test_path",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
36 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
37 "Path to the prediction folder."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
38 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
39 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
40 parser.add_argument("--save_dir",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
41 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
42 "Path to the directory where bounding boxes text files"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
43 " would be saved."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
44 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
45 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
46 parser.add_argument("--run_dir",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
47 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
48 "Path where overlaid images would be saved."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
49 "For example: `RUN_DIR=projectName/results`."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
50 "This should exist."
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
51 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
52 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
53 parser.add_argument("--foldername",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
54 help=("Folder to save overlaid images.\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
55 "For example: FOLDERNAME=batch.\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
56 "This should not exist as a new folder named `batch`\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
57 " will be created in RUN_DIR.\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
58 " If it exists already then, a new folder named `batch1`\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
59 " will be created automatically as it does not overwrite\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
60 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
61 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
62
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
63 # For selecting and loading model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
64 parser.add_argument("--model_name",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
65 help=("Models for task `detect` can be seen here:\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
66 "https://docs.ultralytics.com/tasks/detect/#models \n\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
67 "Models for task `segment` can be seen here:\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
68 "https://docs.ultralytics.com/tasks/segment/#models \n\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
69 " . Use `yolov8n` for `detect` tasks. "
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
70 "For custom model, use `best`"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
71 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
72 default='yolov8n', type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
73 parser.add_argument("--model_path",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
74 help="Full absolute path to the model directory",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
75 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
76 parser.add_argument("--model_format",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
77 help="Format of the YOLO model i.e pt, yaml etc.",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
78 default='pt', type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
79 parser.add_argument("--class_names_file",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
80 help="Path to the text file containing class names.",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
81 type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
82
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
83 # For training the model and prediction
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
84 parser.add_argument("--mode",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
85 help=(
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
86 "detection, segmentation, classification, and pose \n. "
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
87 " Only detection mode available currently i.e. `detect`"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
88 ), default='detect', type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
89 parser.add_argument('--train',
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
90 help="Do training",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
91 action='store_true')
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
92 parser.add_argument("--confidence",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
93 help="Confidence value (0-1) for each detected bounding box",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
94 default=0.5, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
95 parser.add_argument("--epochs",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
96 help="Number of epochs for training. Default: 100",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
97 default=100, type=int)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
98 parser.add_argument("--init_lr",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
99 help="Number of epochs for training. Default: 100",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
100 default=0.01, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
101 parser.add_argument("--weight_decay",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
102 help="Number of epochs for training. Default: 100",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
103 default=0.0005, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
104
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
105 parser.add_argument("--num_classes",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
106 help="Number of classes to be predicted. Default: 2",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
107 default=2, type=int)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
108 parser.add_argument("--iou",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
109 help="Intersection over union (IoU) threshold for NMS",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
110 default=0.7, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
111 parser.add_argument("--image_size",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
112 help=("Size of input image to be used only as integer of w,h. \n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
113 "For training choose <= 1000. \n\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
114 "Prediction will be done on original image size"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
115 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
116 default=320, type=int)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
117 parser.add_argument("--max_det",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
118 help=("Maximum number of detections allowed per image. \n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
119 "Limits the total number of objects the model can detect in a single inference, \n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
120 "preventing excessive outputs in dense scenes.\n\n"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
121 ),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
122 default=300, type=int)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
123
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
124 # For tracking
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
125 parser.add_argument("--tracker_file",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
126 help=("Path to the configuration file of the tracker used. \n"),
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
127 default='bytetrack.yaml', type=str)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
128
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
129 # For headless operation
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
130 parser.add_argument('--headless', action='store_true')
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
131 parser.add_argument('--nextflow', action='store_true')
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
132
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
133 # For data augmentation
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
134 parser.add_argument("--hsv_h",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
135 help="(float) image HSV-Hue augmentation (fraction)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
136 default=0.015, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
137 parser.add_argument("--hsv_s",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
138 help="(float) image HSV-Saturation augmentation (fraction)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
139 default=0.7, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
140 parser.add_argument("--hsv_v",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
141 help="(float) image HSV-Value augmentation (fraction)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
142 default=0.4, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
143 parser.add_argument("--degrees",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
144 help="(float) image rotation (+/- deg)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
145 default=0.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
146 parser.add_argument("--translate",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
147 help="(float) image translation (+/- fraction)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
148 default=0.1, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
149 parser.add_argument("--scale",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
150 help="(float) image scale (+/- gain)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
151 default=0.5, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
152 parser.add_argument("--shear",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
153 help="(float) image shear (+/- deg)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
154 default=0.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
155 parser.add_argument("--perspective",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
156 help="(float) image perspective (+/- fraction), range 0-0.001",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
157 default=0.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
158 parser.add_argument("--flipud",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
159 help="(float) image flip up-down (probability)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
160 default=0.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
161 parser.add_argument("--fliplr",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
162 help="(float) image flip left-right (probability)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
163 default=0.5, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
164 parser.add_argument("--mosaic",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
165 help="(float) image mosaic (probability)",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
166 default=1.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
167 parser.add_argument("--crop_fraction",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
168 help="(float) crops image to a fraction of its size to "
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
169 "emphasize central features and adapt to object scales, "
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
170 "reducing background distractions",
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
171 default=1.0, type=float)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
172
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
173
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
174 #
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
175 # Functions
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
176 #
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
177 # Train a new model on the dataset mentioned in yaml file
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
178 def trainModel(model_path, model_name, yaml_filepath, **kwargs):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
179 if "imgsz" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
180 image_size = kwargs['imgsz']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
181 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
182 image_size = 320
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
183
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
184 if "epochs" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
185 n_epochs = kwargs['epochs']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
186 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
187 n_epochs = 100
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
188
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
189 if "hsv_h" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
190 aug_hsv_h = kwargs['hsv_h']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
191 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
192 aug_hsv_h = 0.015
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
193
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
194 if "hsv_s" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
195 aug_hsv_s = kwargs['hsv_s']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
196 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
197 aug_hsv_s = 0.7
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
198
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
199 if "hsv_v" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
200 aug_hsv_v = kwargs['hsv_v']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
201 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
202 aug_hsv_v = 0.4
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
203
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
204 if "degrees" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
205 aug_degrees = kwargs['degrees']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
206 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
207 aug_degrees = 10.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
208
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
209 if "translate" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
210 aug_translate = kwargs['translate']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
211 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
212 aug_translate = 0.1
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
213
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
214 if "scale" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
215 aug_scale = kwargs['scale']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
216 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
217 aug_scale = 0.2
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
218
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
219 if "shear" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
220 aug_shear = kwargs['shear']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
221 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
222 aug_shear = 0.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
223
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
224 if "shear" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
225 aug_shear = kwargs['shear']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
226 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
227 aug_shear = 0.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
228
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
229 if "perspective" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
230 aug_perspective = kwargs['perspective']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
231 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
232 aug_perspective = 0.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
233
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
234 if "fliplr" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
235 aug_fliplr = kwargs['fliplr']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
236 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
237 aug_fliplr = 0.5
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
238
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
239 if "flipud" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
240 aug_flipud = kwargs['flipud']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
241 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
242 aug_flipud = 0.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
243
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
244 if "mosaic" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
245 aug_mosaic = kwargs['mosaic']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
246 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
247 aug_mosaic = 1.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
248
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
249 if "crop_fraction" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
250 aug_crop_fraction = kwargs['crop_fraction']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
251 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
252 aug_crop_fraction = 1.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
253
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
254 if "weight_decay" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
255 weight_decay = kwargs['weight_decay']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
256 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
257 weight_decay = 1.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
258
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
259 if "init_lr" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
260 init_lr = kwargs['init_lr']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
261 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
262 init_lr = 1.0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
263
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
264 # Load a pretrained YOLO model (recommended for training)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
265 if args.model_format == 'pt':
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
266 model = YOLO(os.path.join(model_path, model_name + "." + args.model_format))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
267 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
268 model = YOLO(model_name + "." + args.model_format)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
269 model.train(data=yaml_filepath, epochs=n_epochs, project=args.run_dir,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
270 imgsz=image_size, verbose=True, hsv_h=aug_hsv_h,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
271 hsv_s=aug_hsv_s, hsv_v=aug_hsv_v, degrees=aug_degrees,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
272 translate=aug_translate, shear=aug_shear, scale=aug_scale,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
273 perspective=aug_perspective, fliplr=aug_fliplr,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
274 flipud=aug_flipud, mosaic=aug_mosaic, crop_fraction=aug_crop_fraction,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
275 weight_decay=weight_decay, lr0=init_lr, seed=42)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
276 return model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
277
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
278
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
279 # Validate the trained model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
280 def validateModel(model):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
281 # Validate the model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
282 metrics = model.val() # no args needed, dataset & settings remembered
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
283 metrics.box.map # map50-95
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
284 metrics.box.map50 # map50
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
285 metrics.box.map75 # map75
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
286 metrics.box.maps # a list contains map50-95 of each category
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
287
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
288
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
289 # Do predictions on images/videos using trained/loaded model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
290 def predict(model, source_datapath, **kwargs):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
291 if "imgsz" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
292 image_size = kwargs['imgsz']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
293 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
294 image_size = 320
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
295
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
296 if "conf" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
297 confidence = kwargs['conf']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
298 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
299 confidence = 0.5
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
300
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
301 if "iou" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
302 iou_value = kwargs['iou']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
303 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
304 iou_value = 0.5
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
305
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
306 if "num_classes" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
307 class_array = list(range(kwargs['num_classes']))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
308 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
309 class_array = [0, 1]
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
310
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
311 if "max_det" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
312 maximum_detections = args.max_det
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
313 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
314 maximum_detections = 300
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
315
3
97bc82ee2a61 planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 743c8acf1ea4e4b1e718743d3772b7e592646611
bgruening
parents: 2
diff changeset
316 run_save_dir = kwargs['run_dir'] # For Galaxy, run_save_dir is always provided via xml wrapper
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
317 if "foldername" in kwargs:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
318 save_folder_name = kwargs['foldername']
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
319 # infer on a local image or directory containing images/videos
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
320 prediction = model.predict(source=source_datapath, save=True, stream=True,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
321 conf=confidence, imgsz=image_size,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
322 save_conf=True, iou=iou_value, max_det=maximum_detections,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
323 classes=class_array, save_txt=False,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
324 project=run_save_dir, name=save_folder_name, verbose=True)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
325 return prediction
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
326
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
327
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
328 # Save bounding boxes
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
329 def save_yolo_bounding_boxes_to_txt(predictions, save_dir):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
330 """
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
331 Function to save YOLO bounding boxes to text files.
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
332 Parameters:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
333 - predictions: List of results from YOLO model inference.
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
334 - save_dir: Directory where the text files will be saved.
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
335 """
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
336 for result in predictions:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
337 result = result.to("cpu").numpy()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
338 # Using bounding_boxes, confidence_scores, and class_num which are defined in the list
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
339 bounding_boxes = result.boxes.xyxy # Bounding boxes in xyxy format
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
340 confidence_scores = result.boxes.conf # Confidence scores
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
341 class_nums = result.boxes.cls # Class numbers
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
342 # Create save directory if it doesn't exist
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
343 save_path = pathlib.Path(save_dir).absolute()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
344 save_path.mkdir(parents=True, exist_ok=True)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
345 # Construct filename for the text file
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
346 image_filename = pathlib.Path(result.path).stem
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
347 text_filename = save_path / f"{image_filename}.txt"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
348 # Write bounding boxes info into the text file
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
349 with open(text_filename, 'w') as f:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
350 for i in range(bounding_boxes.shape[0]):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
351 x1, y1, x2, y2 = bounding_boxes[i]
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
352 confidence = confidence_scores[i]
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
353 class_num = int(class_nums[i])
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
354 f.write(f'{class_num:01} {x1:06.2f} {y1:06.2f} {x2:06.2f} {y2:06.2f} {confidence:0.02} \n')
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
355 print(colored(f"Bounding boxes saved in: {text_filename}", 'green'))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
356
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
357
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
358 if __name__ == '__main__':
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
359 args = parser.parse_args()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
360 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
361 # Train/load model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
362 if (args.train):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
363 model = trainModel(args.model_path, args.model_name, args.yaml_path,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
364 imgsz=args.image_size, epochs=args.epochs,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
365 hsv_h=args.hsv_h, hsv_s=args.hsv_s, hsv_v=args.hsv_v,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
366 degrees=args.degrees, translate=args.translate,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
367 shear=args.shear, scale=args.scale,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
368 perspective=args.perspective, fliplr=args.fliplr,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
369 flipud=args.flipud, mosaic=args.mosaic)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
370 validateModel(model)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
371 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
372 t = time.time()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
373 train_save_path = os.path.expanduser('~/runs/' + args.mode + '/')
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
374 if os.path.isfile(os.path.join(train_save_path,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
375 "train", "weights", "best.pt")) and (args.model_name == 'sam'):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
376 model = YOLO(os.path.join(train_save_path,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
377 "train", "weights", "best.pt"))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
378 else:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
379 model = YOLO(os.path.join(args.model_path,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
380 args.model_name + ".pt"))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
381 model.info(verbose=True)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
382 elapsed = time.time() - t
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
383 print(colored(f"\nYOLO model loaded in : '{elapsed}' sec \n", 'white', 'on_yellow'))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
384
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
385 if (args.save_dir):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
386 # Do predictions (optionally show image results with bounding boxes)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
387 t = time.time()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
388 datapath_for_prediction = args.test_path
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
389 # Extracting class names from the model
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
390 class_names = model.names
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
391 predictions = predict(model, datapath_for_prediction,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
392 imgsz=args.image_size, conf=args.confidence,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
393 iou=args.iou, run_dir=args.run_dir,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
394 foldername=args.foldername, num_classes=args.num_classes, max_det=args.max_det)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
395 elapsed = time.time() - t
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
396 print(colored(f"\nYOLO prediction done in : '{elapsed}' sec \n", 'white', 'on_cyan'))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
397
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
398 if (args.mode == "detect"):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
399 # Save bounding boxes
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
400 save_yolo_bounding_boxes_to_txt(predictions, args.save_dir)
1
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
401
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
402 # Loop over each result
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
403 for result in predictions:
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
404 img = np.copy(result.orig_img)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
405 image_filename = pathlib.Path(result.path).stem
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
406 overlay_path = os.path.join(args.save_dir, f"{image_filename}_overlay.jpg")
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
407
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
408 for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
409 x1, y1, x2, y2 = map(int, box.tolist())
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
410 class_num = int(cls.item())
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
411 confidence = conf.item()
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
412 label = f"{class_names[class_num]} {confidence:.2f}" if class_names else f"{class_num} {confidence:.2f}"
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
413
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
414 cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
415 cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
416 0.5, (0, 255, 0), thickness=1)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
417
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
418 cv2.imwrite(overlay_path, img)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
419 print(colored(f"Overlay image saved at: {overlay_path}", 'cyan'))
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
420 elif (args.mode == "track"):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
421 results = model.track(source=datapath_for_prediction,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
422 tracker=args.tracker_file,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
423 conf=args.confidence,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
424 iou=args.iou,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
425 persist=False,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
426 show=True,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
427 save=True,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
428 project=args.run_dir,
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
429 name=args.foldername)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
430 # Store the track history
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
431 track_history = defaultdict(lambda: [])
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
432
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
433 for result in results:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
434 # Get the boxes and track IDs
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
435 if result.boxes and result.boxes.is_track:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
436 boxes = result.boxes.xywh.cpu()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
437 track_ids = result.boxes.id.int().cpu().tolist()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
438 # Visualize the result on the frame
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
439 frame = result.plot()
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
440 # Plot the tracks
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
441 for box, track_id in zip(boxes, track_ids):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
442 x, y, w, h = box
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
443 track = track_history[track_id]
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
444 track.append((float(x), float(y))) # x, y center point
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
445 if len(track) > 30: # retain 30 tracks for 30 frames
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
446 track.pop(0)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
447
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
448 # Draw the tracking lines
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
449 points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
450 cv2.polylines(frame, [points], isClosed=False, color=(230, 230, 230), thickness=2)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
451
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
452 # Display the annotated frame
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
453 cv2.imshow("YOLO11 Tracking", frame)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
454 print(colored(f"Tracking results saved in : '{args.save_dir}' \n", 'green'))
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
455 elif (args.mode == "segment"):
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
456 # Read class names from the file
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
457 with open(args.class_names_file, 'r') as f:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
458 class_names = [line.strip() for line in f.readlines()]
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
459 class_to_index = {class_name: i for i, class_name in enumerate(class_names)}
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
460
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
461 # Save polygon coordinates
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
462 for result in predictions:
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
463 img = np.copy(result.orig_img)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
464 filename = pathlib.Path(result.path).stem
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
465 b_mask = np.zeros(img.shape[:2], np.uint8)
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
466 mask_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + "_mask.tiff")).absolute())
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
467 txt_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + ".txt")).absolute())
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
468
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
469 for c, ci in enumerate(result):
1
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
470 if ci.masks is not None and ci.masks.xy:
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
471 # Extract contour
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
472 contour = ci.masks.xy.pop()
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
473 contour = contour.astype(np.int32).reshape(-1, 1, 2)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
474 _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
475
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
476 # Normalized polygon points
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
477 points = ci.masks.xyn.pop()
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
478 obj_class = int(ci.boxes.cls.to("cpu").numpy().item())
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
479 confidence = result.boxes.conf.to("cpu").numpy()[c]
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
480
1
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
481 with open(txt_save_as, 'a') as f:
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
482 segmentation_points = ['{} {}'.format(points[i][0], points[i][1]) for i in range(len(points))]
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
483 segmentation_points_string = ' '.join(segmentation_points)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
484 line = '{} {} {}\n'.format(obj_class, segmentation_points_string, confidence)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
485 f.write(line)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
486 else:
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
487 print(colored(f"⚠️ No mask found for object {c} in '{filename}'. Skipping.", "yellow"))
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
488
1
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
489 # Overlay mask onto original image
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
490 colored_mask = cv2.merge([b_mask, np.zeros_like(b_mask), np.zeros_like(b_mask)])
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
491 blended = cv2.addWeighted(img, 1.0, colored_mask, 0.5, 0)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
492 overlay_path = os.path.join(args.save_dir, filename + "_overlay.jpg")
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
493 cv2.imwrite(overlay_path, blended)
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
494
1
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
495 imwrite(mask_save_as, b_mask, imagej=True)
dfda27273ead planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
bgruening
parents: 0
diff changeset
496 print(colored(f"Saved binary mask as : \n '{mask_save_as}' \n", 'magenta'))
0
252fd085940d planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
bgruening
parents:
diff changeset
497 print(colored(f"Polygon coordinates saved as : \n '{txt_save_as}' \n", 'cyan'))