!pip install -q ultralytics
import zipfile, os, shutil, glob
from google.colab import files

# PERBAIKAN (section: unggah dataset & penyesuaian nama otomatis) - dulu
# dataset.zip harus diunggah manual lewat ikon folder di sisi kiri Colab
# DENGAN nama persis "dataset.zip", padahal berkas yang diunduh bisa
# bernama lain tergantung asalnya:
#   - dataset.zip           -> dari tab "Mulai di Sini" (Langkah 1, label di browser)
#   - sampel-dataset.zip     -> dari halaman "Sampel Foto & Label" (tombol Ekspor)
#   - dataset (1).zip, dst   -> browser menambah angka kalau nama bentrok
# Sel ini membuka jendela pilih-berkas langsung di Colab, menerima nama
# APA PUN selama formatnya .zip, lalu mengganti namanya menjadi
# "dataset.zip" secara otomatis supaya sisa kode di bawah selalu
# konsisten tanpa perlu diubah manual tiap kali sumber datasetnya beda.
print("Pilih file dataset .zip (dari 'Mulai di Sini' ATAU dari 'Sampel Foto & Label')...")
uploaded = files.upload()
zip_candidates = [f for f in uploaded.keys() if f.lower().endswith('.zip')]
if not zip_candidates:
    raise RuntimeError("Tidak ada file .zip yang diunggah. Jalankan ulang sel ini lalu pilih file dataset .zip Anda.")
if len(zip_candidates) > 1:
    print(f"Ada {len(zip_candidates)} file .zip diunggah sekaligus - memakai yang pertama: {zip_candidates[0]}")
uploaded_name = zip_candidates[0]
if uploaded_name != 'dataset.zip':
    shutil.move(uploaded_name, 'dataset.zip')
    print(f"Nama berkas '{uploaded_name}' otomatis disesuaikan menjadi 'dataset.zip'.")

if os.path.isdir('dataset'):
    shutil.rmtree('dataset')  # bersihkan sisa ekstraksi dari percobaan sebelumnya kalau ada
zipfile.ZipFile('dataset.zip').extractall('dataset')

# PERBAIKAN: sebagian dataset (mis. hasil zip ulang manual) punya satu
# folder pembungkus tambahan di dalamnya (mis. dataset/dataset/data.yaml)
# alih-alih langsung di akar. Dicari otomatis di mana pun data.yaml
# berada, supaya struktur folder yang sedikit berbeda tetap bisa dipakai
# langsung tanpa perlu dirapikan manual.
yaml_matches = glob.glob('dataset/**/data.yaml', recursive=True)
if not yaml_matches:
    raise RuntimeError("data.yaml tidak ditemukan di dalam dataset.zip. Pastikan file yang diunggah adalah dataset.zip hasil Langkah 1 (bukan berkas lain).")
dataset_dir = os.path.abspath(os.path.dirname(yaml_matches[0]))
yaml_path = os.path.join(dataset_dir, 'data.yaml')

# PERBAIKAN: beberapa versi Ultralytics salah mengartikan "path: ." di
# data.yaml (mencari dataset di folder default ~/datasets/ Colab,
# BUKAN folder hasil extract di atas) - ini penyebab paling umum error
# "dataset tidak ditemukan/kosong". Diperbaiki di sini dengan mengganti
# "path: ." menjadi path folder dataset yang SEBENARNYA (absolut).
with open(yaml_path) as f:
    content = f.read().replace('path: .', f'path: {dataset_dir}')
with open(yaml_path, 'w') as f:
    f.write(content)
print(open(yaml_path).read())  # tampilkan isi data.yaml final untuk dicek sekilas

from ultralytics import YOLO
model = YOLO('yolov8n.pt')          # model dasar kecil & cepat, cocok untuk pemula
model.train(data=yaml_path, epochs=80, imgsz=640)

# Ekspor ke .onnx - imgsz & opset harus SAMA seperti di bawah ini supaya
# kompatibel dengan ONNX Runtime yang dipakai AI Server (lihat
# scripts/export_onnx.py di server untuk penjelasan lengkap).
model.export(format='onnx', imgsz=640, opset=12, simplify=True)
print("Selesai! File .onnx ada di runs/detect/train/weights/best.onnx - unduh lewat panel folder di kiri (klik kanan > Download).")
print("PENTING: saat unggah di Langkah 3 di bawah, isi kolom Label dengan PERSIS 8 label berikut, urutannya JANGAN diubah:")
print('["Sampah","Vegetasi","Genangan Air","Kotoran","Area Bersih","Kerapihan","Kondisi Umum","Berantakan"]')