54 lines
2.5 KiB
Python
54 lines
2.5 KiB
Python
import os
|
|
import glob
|
|
import shutil
|
|
from ultralytics import YOLO
|
|
|
|
# 1. Nạp mô hình PyTorch gốc từ file bạn đã tải sẵn
|
|
model = YOLO("yolov11n-seg.pt")
|
|
|
|
print("⏳ Bước 1: Đang dịch sang định dạng trung gian ONNX...")
|
|
# Xuất sang file .onnx (Vượt qua bộ check OS của LiteRT)
|
|
onnx_path = model.export(format="onnx", imgsz=640)
|
|
|
|
print("⏳ Bước 2: Kích hoạt onnx2tf để dịch ma trận đồ họa sang TFLite...")
|
|
onnx_file = "yolov11n-seg.onnx"
|
|
|
|
if os.path.exists(onnx_file):
|
|
# Gọi công cụ dòng lệnh chuyển dịch sang Float32
|
|
# -nuo: Chống tối ưu hóa lỗi cấu hình phần cứng Android GPU Delegate
|
|
os.system(f"onnx2tf -in {onnx_file} -nuo")
|
|
|
|
print("🔍 Đang quét tìm vị trí file .tflite thành phẩm...")
|
|
# 🟢 SỬA LỖI: Tìm kiếm tất cả các file .tflite được sinh ra trong các thư mục con
|
|
tflite_matches = glob.glob("**/yolov11n-seg*.tflite", recursive=True) + glob.glob("*.tflite")
|
|
|
|
# Lọc bỏ nếu tìm trúng file đích cũ để tránh trùng lặp
|
|
tflite_matches = [f for f in tflite_matches if "yolov11n_seg_portrait.tflite" not in f]
|
|
|
|
if tflite_matches:
|
|
generated_tflite = tflite_matches[0]
|
|
dest_path = "yolov11n_seg_portrait.tflite"
|
|
|
|
# Di chuyển và đổi tên chuẩn xác cho dự án Android
|
|
os.replace(generated_tflite, dest_path)
|
|
print(f"🎉 XUẤT FILE THÀNH CÔNG! File đã nằm tại: {os.path.abspath(dest_path)}")
|
|
|
|
# --- DỌN DẸP FILE RÁC TRUNG GIAN ---
|
|
try:
|
|
os.remove(onnx_file)
|
|
# Xóa các thư mục tạm do onnx2tf hoặc ultralytics tạo ra
|
|
if os.path.exists("yolov11n-seg_saved_model"):
|
|
shutil.rmtree("yolov11n-seg_saved_model")
|
|
if os.path.exists("saved_model"):
|
|
shutil.rmtree("saved_model")
|
|
# Tìm và xóa thư mục con trùng tên nếu có
|
|
for d in os.listdir("."):
|
|
if os.path.isdir(d) and "yolov11n-seg" in d:
|
|
shutil.rmtree(d)
|
|
except Exception as clean_ex:
|
|
print(f"⚠️ Cảnh báo dọn dẹp: {str(clean_ex)}")
|
|
|
|
else:
|
|
print("❌ Lỗi: onnx2tf không tạo ra bất kỳ file .tflite nào! Hãy kiểm tra log phía trên của onnx2tf xem có bị lỗi thiếu công cụ 'flatc' không.")
|
|
else:
|
|
print("❌ Lỗi: Không tìm thấy file ONNX trung gian để chuyển đổi!") |