agrobot_base/Python/OAK/datasets/download_backbone.py

12 lines
386 B
Python

from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
modelo = "nvidia/mit-b0"
pasta = r"C:\AgroBaseModels\Backbones\nvidia_mit_b0"
model = SegformerForSemanticSegmentation.from_pretrained(modelo)
processor = SegformerImageProcessor.from_pretrained(modelo)
model.save_pretrained(pasta)
processor.save_pretrained(pasta)
print("Backbone salvo em:", pasta)