Coverage for drawyolo/main.py: 100.00%

1from pathlib import Path 

2import typer 

3 

4from .draw import draw_box_on_image_with_labels, draw_box_on_image_with_model 

5 

6app = typer.Typer() 

7 

8 

9@app.command() 

10def drawyolo( 

11 image: Path = typer.Argument(..., help="Input image path"), 

12 output: Path = typer.Argument(..., help="Output image path"), 

13 labels: Path = typer.Option(None, help="Path to labels text file in YOLO format"), 

14 weights: Path = typer.Option( 

15 None, help="Path to YOLO model weights file if not using labels text file" 

16 ), 

17 classes: str = typer.Option(None, help="Class names separated by comma"), 

18 highest: bool = typer.Option( 

19 False, help="Only draw the box with the highest confidence" 

20 ), 

21 width: int = typer.Option( 

22 None, help="Width of the output image if you wish to resize it" 

23 ), 

24 height: int = typer.Option( 

25 None, help="Height of the output image if you wish to resize it" 

26 ), 

27 line_thickness: int = typer.Option( 

28 None, 

29 help="Thickness of the box lines. If not provided, it will be calculated based on the final image size.", 

30 ), 

31): 

32 """ 

33 Draw boxes on image from a YOLO model or labels text file in YOLO format. 

34 """ 

35 if classes is not None: 

36 classes = classes.split(",") 

37 

38 if labels is not None: 

39 draw_box_on_image_with_labels( 

40 image, 

41 labels, 

42 output, 

43 classes, 

44 width=width, 

45 height=height, 

46 line_thickness=line_thickness, 

47 ) 

48 elif weights is not None: 

49 draw_box_on_image_with_model( 

50 image, 

51 weights, 

52 output, 

53 classes=classes, 

54 highest=highest, 

55 width=width, 

56 height=height, 

57 line_thickness=line_thickness, 

58 ) 

59 else: 

60 raise typer.BadParameter("You must provide either --labels or --weights.")