Coverage for scripts/training.py: 65%

1from ultralytics import YOLO 

2import shutil 

3import os 

4import torch 

5import random 

6from collections import defaultdict 

7 

8def copy_files(source_dir, dest_dir): 

9 os.makedirs(dest_dir, exist_ok=True) 

10 # Iterate through all files in source directory 

11 for item in os.listdir(source_dir): 

12 if item == "old_models": 

13 continue 

14 source_path = os.path.join(source_dir, item) 

15 dest_path = os.path.join(dest_dir, item) 

16 print(f"Copying {item} from {source_path} to {dest_path}") 

17 shutil.copytree(source_path, dest_path, dirs_exist_ok=True) 

18 

19source_directory = "../models/" 

20destination_directory = "../models/old_models/" 

21# copy_files(source_directory, destination_directory) 

22# shutil.rmtree("../models/second_training/", ignore_errors=True) 

23# shutil.rmtree("../models/yolo-object-lane-unfroze/", ignore_errors=True) 

24 

25def train_objects_model(): 

26 model = YOLO("yolov8n-seg.pt") 

27 return model.train( 

28 data="/home/seame/ObjectDetectionAvoidance/yolo_models/split_dataset/data.yaml", 

29 epochs=150, 

30 warmup_epochs=5, 

31 imgsz=320, 

32 hsv_h=0.4, #hue 

33 hsv_s=0.7, # saturation 

34 hsv_v=0.4, #brightness 

35 translate=0.3, # Moderate translation 

36 scale=0.4, 

37 batch=16, 

38 device=0, 

39 workers=8, 

40 project="../models", 

41 name="objects", 

42 exist_ok=True, 

43 freeze=None, # Unfreeze all layers 

44 lr0=0.002, 

45 patience=20, # Early stopping 

46 weight_decay=0.001, 

47 fliplr=0, # horizontal flip 

48 cls=1.5, # Emphasize classification loss 

49 box=7.5, # Default 

50 dfl=1.5, 

51 label_smoothing=0.1, 

52 mosaic=0.05, 

53 erasing=0.5, 

54 mixup=0.4, # Add mixup for small dataset 

55 copy_paste=0.5, 

56 auto_augment=None, # Disable auto-augmentation 

57 ) 

58 

59 

60def train_seame_model(): 

61 model = YOLO("../models/objects/weights/best.pt") # Load the trained model 

62 return model.train( 

63 data="/home/seame/ObjectDetectionAvoidance/yolo_models/seame_training/data.yaml", 

64 epochs=150, 

65 warmup_epochs=5, 

66 imgsz=512, 

67 hsv_h=0.5, #hue 

68 hsv_s=0.7, # saturation 

69 hsv_v=0.4, #brightness 

70 translate=0.4, # Moderate translation 

71 scale=0.5, 

72 batch=16, 

73 device=0, 

74 workers=8, 

75 project="../models", 

76 name="seame_n", 

77 exist_ok=True, 

78 freeze=0, # Unfreeze all layers 

79 lr0=0.002, 

80 patience=20, # Early stopping 

81 weight_decay=0.001, 

82 fliplr=0, # horizontal flip 

83 mosaic=0, 

84 erasing=0.2, 

85 cls=1.5, # Emphasize classification loss 

86 box=7.5, # Default 

87 dfl=1.5, 

88 label_smoothing=0.15, 

89 mixup=0, 

90 copy_paste=0, 

91 auto_augment=None, # Disable auto-augmentation 

92 )