be/main.py

import base64
import io
import os
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
from dotenv import load_dotenv
from groq import Groq

load_dotenv()
# Trigger reload to catch new .env key

# Setup Groq API Key
# Pastikan Anda menambahkan variabel GROQ_API_KEY di file .env Anda
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY", "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"))

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load the model (YOLOv26x untuk akurasi maksimal seperti yang diminta)
try:
    model = YOLO('../model/waste_segmentor_yolo26/waste_yolo26seg_best.pt')
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

# Daftar nama kelas COCO yang termasuk E-Waste (elektronik) — digunakan untuk override LLM
E_WASTE_CLASSES = {
    "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
    "microwave", "oven", "toaster", "refrigerator", "hair drier"
}

# Filter khusus untuk COCO dataset ke 4 jenis sampah
coco_to_waste = {
    "bottle": {"cat": "Anorganik", "color": "#eab308", "desc": "Botol plastik atau kaca yang bisa didaur ulang."},
    "wine glass": {"cat": "Anorganik", "color": "#eab308", "desc": "Pecahan kaca bisa didaur ulang."},
    "cup": {"cat": "Anorganik", "color": "#eab308", "desc": "Gelas plastik atau kertas."},
    "fork": {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang jika logam/plastik."},
    "knife": {"cat": "Anorganik", "color": "#eab308", "desc": "Logam, kumpulkan di bank sampah."},
    "spoon": {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang."},
    "bowl": {"cat": "Anorganik", "color": "#eab308", "desc": "Mangkuk bisa didaur ulang."},
    "banana": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa buah, bisa dijadikan kompos."},
    "apple": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa buah, bisa dijadikan kompos."},
    "sandwich": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa makanan, bisa dijadikan kompos."},
    "orange": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa buah, bisa dijadikan kompos."},
    "broccoli": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa sayur, bisa dijadikan kompos."},
    "carrot": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa sayur, bisa dijadikan kompos."},
    "hot dog": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa makanan basah."},
    "pizza": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa makanan basah."},
    "donut": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa makanan basah."},
    "cake": {"cat": "Organik", "color": "#16a34a", "desc": "Sisa makanan basah."},
    "tv": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste, mengandung komponen berbahaya."},
    "laptop": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste, mengandung logam berat dan baterai."},
    "mouse": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste elektronik."},
    "remote": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste elektronik."},
    "keyboard": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste elektronik."},
    "cell phone": {"cat": "B3", "color": "#dc2626", "desc": "E-Waste, mengandung baterai li-ion."},
    "microwave": {"cat": "B3", "color": "#dc2626", "desc": "Limbah elektronik besar."},
    "oven": {"cat": "B3", "color": "#dc2626", "desc": "Limbah elektronik besar."},
    "toaster": {"cat": "B3", "color": "#dc2626", "desc": "Limbah elektronik."},
    "refrigerator": {"cat": "B3", "color": "#dc2626", "desc": "Limbah elektronik besar (freon berbahaya)."},
    "backpack": {"cat": "Residu", "color": "#475569", "desc": "Bahan kain campuran sulit didaur ulang."},
    "umbrella": {"cat": "Residu", "color": "#475569", "desc": "Bahan campuran sulit didaur."},
    "handbag": {"cat": "Residu", "color": "#475569", "desc": "Bahan campuran kulit/kain sulit didaur."},
    "tie": {"cat": "Residu", "color": "#475569", "desc": "Kain tekstil."},
    "suitcase": {"cat": "Residu", "color": "#475569", "desc": "Bahan keras campuran."},
    "book": {"cat": "Anorganik", "color": "#eab308", "desc": "Kertas bisa didaur ulang."},
    "vase": {"cat": "Residu", "color": "#475569", "desc": "Keramik sulit didaur ulang."},
    "scissors": {"cat": "Anorganik", "color": "#eab308", "desc": "Logam tajam."},
    "teddy bear": {"cat": "Residu", "color": "#475569", "desc": "Bahan kain dan busa."},
    "hair drier": {"cat": "B3", "color": "#dc2626", "desc": "E-waste elektronik."},
    "toothbrush": {"cat": "Residu", "color": "#475569", "desc": "Plastik keras bercampur sikat, sulit didaur."}
}

# Categories mapping based on data.yaml (fallback jika menggunakan model custom)
class_categories = {
    0: {"cat": "B3", "color": "#ef4444", "desc": "Limbah Bahan Berbahaya dan Beracun."},
    1: {"cat": "B3", "color": "#ef4444", "desc": "Limbah elektronik yang mengandung logam berat."},
    2: {"cat": "B3", "color": "#ef4444", "desc": "Sisa cat yang berbahaya bagi lingkungan."},
    3: {"cat": "B3", "color": "#ef4444", "desc": "Bahan kimia pembasmi hama."},
    4: {"cat": "Residu", "color": "#6b7280", "desc": "Sulit didaur ulang dan tidak dapat dikomposkan."},
    5: {"cat": "Residu", "color": "#6b7280", "desc": "Sulit didaur ulang dan tidak dapat dikomposkan."},
    6: {"cat": "Residu", "color": "#6b7280", "desc": "Plastik lembaran atau pembungkus yang sulit didaur ulang."},
    7: {"cat": "Residu", "color": "#6b7280", "desc": "Sulit didaur ulang dan tidak dapat dikomposkan."},
    8: {"cat": "Residu", "color": "#6b7280", "desc": "Styrofoam tidak dapat terurai secara alami."},
    9: {"cat": "Organik", "color": "#22c55e", "desc": "Bisa dijadikan kompos (green material)."},
    10: {"cat": "Organik", "color": "#22c55e", "desc": "Bisa dijadikan kompos untuk kalsium tanah."},
    11: {"cat": "Organik", "color": "#22c55e", "desc": "Bisa dijadikan kompos (green material)."},
    12: {"cat": "Organik", "color": "#22c55e", "desc": "Bisa dijadikan kompos."},
    13: {"cat": "Organik", "color": "#22c55e", "desc": "Sisa potongan tanaman untuk kompos (brown/green material)."},
    14: {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang, kumpulkan di bank sampah."},
    15: {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang, kumpulkan di bank sampah."},
    16: {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang, pastikan tetap kering."},
    17: {"cat": "Anorganik", "color": "#eab308", "desc": "Bisa didaur ulang, kumpulkan di bank sampah."}
}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if not model:
        return {"error": "Model not loaded properly. Please check server logs."}
        
    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            return {"error": "Invalid image format or image could not be read."}
        
        # --- OPTIMASI KECEPATAN (SPEED UP) ---
        # 1. Resize gambar jika terlalu besar agar memori ringan dan upload ke Gemini instan (max 640px)
        max_dim = 640
        h, w = img.shape[:2]
        if max(h, w) > max_dim:
            scale = max_dim / max(h, w)
            img = cv2.resize(img, (int(w * scale), int(h * scale)))

        # Run YOLO prediction
        try:
            # 2. Paksa YOLO menggunakan ukuran gambar lebih kecil (imgsz=320) agar inferensi di CPU laptop sangat cepat
            # Menurunkan dari 640 ke 320 memotong beban komputasi YOLO hingga 75%
            results = model(img, imgsz=320)
        except Exception as model_err:
            print(f"YOLO Model Inference Error: {str(model_err)}")
            return {"error": f"AI model inference failed: {str(model_err)}"}
        
        # Process results
        predictions = []
        yolo_proposals = []
        annotated_img = img.copy()
        detected_items = set()
        
        if results is not None:
            for r in results:
                # Safely parse bounding boxes
                if hasattr(r, 'boxes') and r.boxes is not None:
                    for box in r.boxes:
                        try:
                            class_id = int(box.cls[0])
                            conf = float(box.conf[0])
                            
                            # Safely get class name
                            class_name = "Unknown"
                            if hasattr(model, 'names') and class_id in model.names:
                                class_name = model.names[class_id]
                            
                            # Jika menggunakan yolov26x (COCO), saring hanya barang sampah
                            if len(model.names) > 20: 
                                if class_name not in coco_to_waste:
                                    continue # Abaikan "person", "car", "dog", dll.
                                cat_info = coco_to_waste[class_name]
                            else:
                                cat_info = class_categories.get(class_id, {"cat": "Unknown", "color": "#9ca3af", "desc": "Tidak diketahui."})
                            
                            detected_items.add(class_name)
                            
                            # Simpan deteksi YOLO sebagai proposal (belum digambar ke gambar)
                            yolo_proposals.append({
                                "box": tuple(map(int, box.xyxy[0])),
                                "class_id": class_id,
                                "class_name": class_name,
                                "confidence": conf,
                                "category": cat_info["cat"],
                                "raw_class_name": class_name,
                                "color": cat_info["color"],
                                "description": cat_info["desc"]
                            })
                        except Exception as box_err:
                            print(f"Error parsing detection box: {str(box_err)}")
                            continue
                
        # --- TAHAP 1.5: BYPASS LLM JIKA YOLO SANGAT YAKIN ---
        bypass_llm = False
        if len(yolo_proposals) > 0:
            highest_conf = max([p["confidence"] for p in yolo_proposals])
            if highest_conf >= 0.50: # Threshold tinggi agar kasus bias (rokok/wortel) tetap butuh validasi
                bypass_llm = True

        # --- TAHAP 2: VALIDASI KECERDASAN BUATAN (GROQ LLAMA VISION) ---
        gemini_analysis = "Analisis AI belum siap."
        benda_terdeteksi = "Objek Tak Dikenal"
        fallback_cat = None
        fallback_color = "#9ca3af"
        status_yolo = "SALAH"

        try:
            if bypass_llm:
                status_yolo = "BENAR"
                gemini_analysis = "Status YOLO: BENAR\nBypass API: YOLO memiliki tingkat keyakinan sangat tinggi (>50%)."
            else:
                # Konversi gambar ke Base64 (Syarat wajib untuk Groq API)
                _, img_encoded = cv2.imencode('.jpg', img)
                base64_image = base64.b64encode(img_encoded).decode('utf-8')
                
                yolo_result_str = ", ".join(detected_items) if detected_items else "Tidak ada objek yang terdeteksi dengan jelas"
                
                prompt = (
                    f"Sistem deteksi awal (YOLO) mendeteksi objek ini sebagai: '{yolo_result_str}'. "
                    "PERINGATAN KRITIS: Jangan langsung percaya pada hasil YOLO! YOLO sering salah menebak bentuk (contoh: puntung rokok dikira wortel/carrot, botol dikira vas). "
                    "Gunakan penglihatan visualmu sendiri. Perhatikan tekstur, abu, dan bentuk aslinya. "
                    "ATURAN WAJIB KATEGORI B3 (E-WASTE): Semua barang elektronik WAJIB masuk kategori B3, termasuk tapi tidak terbatas pada: "
                    "TV/monitor/layar, HP/smartphone/handphone, laptop, keyboard, mouse, remote, microwave, oven, toaster, kulkas/refrigerator, hair dryer, "
                    "dan SEMUA perangkat elektronik lainnya. Barang elektronik TIDAK BOLEH masuk ANORGANIK. "
                    "Berikan analisis akhir dengan format teks persis seperti di bawah ini:\n"
                    "Status YOLO: [Isi dengan BENAR jika tebakan YOLO 100% akurat, atau SALAH jika YOLO bias/salah tebak]\n"
                    "Nama Benda: [Isi dengan nama asli benda/sampah di gambar, misal: Puntung Rokok, Botol Plastik, dll, max 5 kata]\n"
                    "Tips: [Isi dengan 1 tips cerdas cara membuang/mendaur ulangnya]\n"
                    "Kategori: [Pilih SATU saja: ORGANIK, ANORGANIK, B3, RESIDU. Catatan: Puntung rokok/popok/tisu kotor wajib masuk RESIDU. Semua elektronik WAJIB B3]"
                )
                
                # Memanggil API Groq (Llama 4 Vision Scout)
                response = groq_client.chat.completions.create(
                    model="meta-llama/llama-4-scout-17b-16e-instruct",
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {"type": "text", "text": prompt},
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/jpeg;base64,{base64_image}",
                                    }
                                }
                            ]
                        }
                    ],
                    temperature=0.4,
                    max_tokens=1024,
                )
                
                gemini_analysis = response.choices[0].message.content.strip()
                
                # Parsing format respons secara cerdas
                tips_cerdas = "Daur ulang atau buang ke tempat sampah yang sesuai."
                for line in gemini_analysis.split('\n'):
                    line = line.strip()
                    line_lower = line.lower()
                    
                    if line_lower.startswith("status yolo:") or line_lower.startswith("**status yolo"):
                        val = line.split(":", 1)[1].strip().upper().replace("*", "")
                        if "BENAR" in val:
                            status_yolo = "BENAR"
                    elif line_lower.startswith("nama benda:") or line_lower.startswith("**nama benda"):
                        parts = line.split(":", 1)
                        if len(parts) > 1:
                            benda_terdeteksi = parts[1].strip().replace("*", "")
                    elif line_lower.startswith("tips:") or line_lower.startswith("**tips"):
                        parts = line.split(":", 1)
                        if len(parts) > 1:
                            tips_cerdas = parts[1].strip().replace("*", "")
                    elif line_lower.startswith("kategori:") or line_lower.startswith("**kategori"):
                        parts = line.split(":", 1)
                        if len(parts) > 1:
                            cat_str = parts[1].strip().upper().replace("*", "")
                            if "B3" in cat_str:
                                fallback_cat = "B3"
                                fallback_color = "#dc2626"
                            elif "ANORGANIK" in cat_str:
                                fallback_cat = "Anorganik"
                                fallback_color = "#eab308"
                            elif "ORGANIK" in cat_str:
                                fallback_cat = "Organik"
                                fallback_color = "#16a34a"
                            elif "RESIDU" in cat_str:
                                fallback_cat = "Residu"
                                fallback_color = "#475569"
            
            # Fallback agresif jika regex kategori gagal
            if fallback_cat is None:
                text_upper = gemini_analysis.upper()
                if "B3" in text_upper:
                    fallback_cat = "B3"
                    fallback_color = "#dc2626"
                elif "ANORGANIK" in text_upper:
                    fallback_cat = "Anorganik"
                    fallback_color = "#eab308"
                elif "ORGANIK" in text_upper:
                    fallback_cat = "Organik"
                    fallback_color = "#16a34a"
                elif "RESIDU" in text_upper:
                    fallback_cat = "Residu"
                    fallback_color = "#475569"

            # POST-PROCESSING: Override LLM jika YOLO mendeteksi barang elektronik
            # LLM kadang salah mengkategorikan elektronik sebagai Anorganik
            yolo_detected_ewaste = detected_items & E_WASTE_CLASSES
            if yolo_detected_ewaste and fallback_cat != "B3":
                print(f"[E-WASTE OVERRIDE] YOLO detected e-waste items: {yolo_detected_ewaste}, but LLM said '{fallback_cat}'. Forcing B3.")
                fallback_cat = "B3"
                fallback_color = "#dc2626"
                    
            # --- DECISION LOGIC: YOLO vs GEMINI ---
            if status_yolo == "BENAR" and len(yolo_proposals) > 0:
                # YOLO benar, gambar bounding box dan gunakan hasil YOLO ke UI
                for prop in yolo_proposals:
                    hex_color = prop["color"].lstrip('#')
                    bgr_color = tuple(int(hex_color[i:i+2], 16) for i in (4, 2, 0))
                    x1, y1, x2, y2 = prop["box"]
                    
                    # Gambar kotak
                    cv2.rectangle(annotated_img, (x1, y1), (x2, y2), bgr_color, 3)
                    
                    # Gambar label yang cerdas
                    label = f"{prop['raw_class_name'].title()} ({prop['category']}) {prop['confidence']:.2f}"
                    (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
                    cv2.rectangle(annotated_img, (x1, y1 - 25), (x1 + w, y1), bgr_color, -1)
                    cv2.putText(annotated_img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

                    predictions.append({
                        "class_id": prop["class_id"],
                        "class_name": prop["class_name"],
                        "confidence": prop["confidence"],
                        "category": prop["category"],
                        "color": prop["color"],
                        "description": prop["description"]
                    })
            else:
                # YOLO salah (bias) atau tidak mendeteksi apa-apa, batalkan bounding box
                # Gunakan hasil murni dari Gemini
                if not fallback_cat:
                    fallback_cat = "Anorganik"
                    fallback_color = "#eab308"

                predictions.append({
                    "class_id": 99,
                    "class_name": benda_terdeteksi.title(),
                    "confidence": 0.99,
                    "category": fallback_cat,
                    "color": fallback_color,
                    "description": tips_cerdas
                })
        except Exception as gemini_err:
            print(f"Groq API Error: {str(gemini_err)}")
            gemini_analysis = "Gagal memanggil API Groq. Pastikan GROQ_API_KEY sudah dimasukkan di file .env"
            # Fallback darurat jika Gemini down, paksakan pakai YOLO
            for prop in yolo_proposals:
                predictions.append({
                    "class_id": prop["class_id"],
                    "class_name": prop["class_name"],
                    "confidence": prop["confidence"],
                    "category": prop["category"],
                    "color": prop["color"],
                    "description": prop["description"]
                })
                
            # Jika YOLO kosong dan Gemini error (misal limit 429), cegah UI Error
            if len(predictions) == 0:
                is_rate_limit = "429" in str(gemini_err) or "quota" in str(gemini_err).lower()
                error_desc = "API Groq melampaui limit gratis (429). Mohon ganti API Key di .env" if is_rate_limit else f"Sistem terganggu: {str(gemini_err)[:50]}"
                predictions.append({
                    "class_id": 99,
                    "class_name": "Deteksi Gagal (Limit API)",
                    "confidence": 0.50,
                    "category": "Residu", # Masuk ke abu-abu jika error
                    "color": "#475569",
                    "description": error_desc
                })
        
        # Convert annotated image to base64 (WAJIB dilakukan di akhir setelah validasi)
        _, buffer = cv2.imencode('.jpg', annotated_img)
        img_base64 = base64.b64encode(buffer).decode('utf-8')

        return {
            "predictions": predictions,
            "image_base64": img_base64,
            "gemini_analysis": gemini_analysis
        }
    except Exception as e:
        print(f"Error during prediction: {str(e)}")
        return {"error": f"Internal server error processing image: {str(e)}"}

@app.get("/")
def read_root():
    return {"message": "KomposIn ML API is running"}