forked from opendatalab/DocLayout-YOLO
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathval.py
More file actions
21 lines (17 loc) · 677 Bytes
/
Copy pathval.py
File metadata and controls
21 lines (17 loc) · 677 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import argparse
from doclayout_yolo import YOLOv10
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--data', default=None, required=True, type=str)
parser.add_argument('--model', default=None, required=True, type=str)
parser.add_argument('--batch-size', default=16, required=False, type=int)
parser.add_argument('--device', default="0,1,2,3,4,5,6,7", required=False, type=str)
args = parser.parse_args()
# Load a pre-trained model
model = YOLOv10(args.model)
# Train the model
model.val(
data=f'{args.data}.yaml',
batch=args.batch_size,
device=args.device,
)