79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
# train.py
|
|
from pathlib import Path
|
|
from ultralytics.models.yolo.segment.train import SegmentationTrainer
|
|
from custom_channel_dataset import CustomChannelSegDataset
|
|
from patch_model_xch import patch_yolov8_first_conv_to_xch
|
|
import yaml
|
|
|
|
with Path("data.yaml").open("r", encoding="utf-8") as f:
|
|
cfg = yaml.safe_load(f) or {}
|
|
CHANNELS = int(cfg.get("channels", 4))
|
|
SCALE = cfg.get("scale", "m")
|
|
IMG_SZ = int(cfg.get("size", 800))
|
|
|
|
class CustomChannelSegTrainer(SegmentationTrainer):
|
|
def build_dataset(self, img_path, mode="train", batch=None):
|
|
# img_path é train.txt ou val.txt vindo do data.yaml
|
|
return CustomChannelSegDataset(
|
|
img_path=img_path,
|
|
data=self.data,
|
|
rect=False,
|
|
canais=CHANNELS,
|
|
imgsz=self.args.imgsz,
|
|
augment=(mode == "train"), # 👈 importante
|
|
hyp=self.args, # 👈 importante: passa os hypers
|
|
)
|
|
|
|
def get_model(self, cfg=None, weights=None, verbose=True):
|
|
model = super().get_model(cfg=cfg, weights=weights, verbose=verbose)
|
|
|
|
patch_yolov8_first_conv_to_xch(model, canais=CHANNELS)
|
|
|
|
# debug pra garantir
|
|
seq = model.model if hasattr(model, "model") else model
|
|
print("First conv in_channels:", seq[0].conv.in_channels)
|
|
|
|
return model
|
|
|
|
|
|
def main():
|
|
args = {
|
|
"project": f"C:/ZendionInc/agrobot_base/Python/yolov8-seg/runs/segment/ch_{CHANNELS}/sc_{SCALE}/sz_{IMG_SZ}",
|
|
"name": "weed_detector_segformer",
|
|
"model": f"backbones/yolov8-seg-{CHANNELS}ch.yaml",
|
|
#"model": f"backbones/yolov8{SCALE}-seg.pt",
|
|
"pretrained": False,
|
|
"data": "data.yaml",
|
|
"epochs": 1000,
|
|
"patience": 200,
|
|
"imgsz": IMG_SZ,
|
|
"device": 0,
|
|
"batch": 4,
|
|
"workers": 4,
|
|
|
|
# zerando todas as augmentações “perigosas” pro debug
|
|
"mosaic": 0.0,
|
|
"copy_paste": 0.0,
|
|
"mixup": 0.0,
|
|
"erasing": 0.0,
|
|
"auto_augment": "none",
|
|
"hsv_h": 0.0,
|
|
"hsv_s": 0.0,
|
|
"hsv_v": 0.0,
|
|
# augmentacoes geometricas
|
|
"fliplr": 0.5,
|
|
"flipud": 0.0,
|
|
"translate": 0.05,
|
|
"scale": 0.2,
|
|
"degrees": 10.0,
|
|
"shear": 0.0,
|
|
"perspective": 0.0
|
|
}
|
|
|
|
trainer = CustomChannelSegTrainer(overrides=args)
|
|
trainer.train()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|