import torch import cv2 import numpy as np import joblib import os import uvicorn from fastapi import FastAPI, UploadFile, File from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware from io import BytesIO from PIL import Image from transformers import AutoImageProcessor, AutoModel from ultralytics import YOLO from sklearn.preprocessing import MultiLabelBinarizer import asyncio # ====================================================== # 0. KONFIGURASI DAN MUAT MODEL (PATH HARUS RELATIF) # ====================================================== # Catatan: Hugging Face Spaces akan mencari file di direktori ini (BASE_PATH = .) BASE_PATH = "." YOLO_MODEL_PATH = os.path.join(BASE_PATH, "best.pt") # Hapus penamaan tanggal yang panjang di Colab dan gunakan nama folder sebenarnya: MODEL_DIR = os.path.join(BASE_PATH, "OVR_Checkpoints-20251018T053026Z-1-001", "OVR_Checkpoints") ENCODER_PATH = os.path.join(BASE_PATH, "dinov3_multilabel_encoder.pkl") MAPPING_SAVE_PATH = os.path.join(BASE_PATH, "label_mapping_dict.joblib") label_columns = ["product", "grade", "cap", "label", "brand", "type", "subtype", "volume"] device = torch.device("cpu") # Gunakan CPU untuk stabilitas di HF Spaces (kecuali Anda memilih GPU di settings) # --- SETUP DINOv3 Auth Token --- # Membaca token dari Environment Variable yang Anda set di Secrets HF Space HF_AUTH_TOKEN = os.environ.get("HUGGINGFACE_TOKEN") # --- FUNGSI EKSTRAKSI DINOv3 (Harus didefinisikan sebelum loading) --- def extract_dinov3_features(image_crop): inputs = dinov3_processor(images=image_crop, return_tensors="pt").to(device) with torch.no_grad(): outputs = dinov3_model(**inputs) cls_token_features = outputs.last_hidden_state[:, 0, :] return cls_token_features.cpu().numpy().flatten() # === LOAD SEMUA MODEL KE MEMORI === try: # 1. Load YOLO yolo_model = YOLO(YOLO_MODEL_PATH) # 2. Load DINOv3 Processor dan Model (Wajib menggunakan token) dinov3_model_name = "facebook/dinov3-convnext-small-pretrain-lvd1689m" dinov3_processor = AutoImageProcessor.from_pretrained(dinov3_model_name, token=HF_AUTH_TOKEN) dinov3_model = AutoModel.from_pretrained(dinov3_model_name, token=HF_AUTH_TOKEN).to(device).eval() # 3. Load Scikit-learn Assets mlb = joblib.load(ENCODER_PATH) mapping_dict = joblib.load(MAPPING_SAVE_PATH) num_classes = len(mlb.classes_) all_classifiers = [] for i in range(num_classes): class_name = mlb.classes_[i] safe_class_name = str(class_name).replace(' ', '_').replace('/', '_').replace(':', '_').replace('.', '_') classifier_path = os.path.join(MODEL_DIR, f"clf_{i}_{safe_class_name}.pkl") if not os.path.exists(classifier_path): classifier_path = os.path.join(MODEL_DIR, f"clf_{i}_{class_name}.pkl") all_classifiers.append(joblib.load(classifier_path)) print("✅ Semua model dan aset dimuat ke memori.") except Exception as e: # Jika gagal memuat, cetak error dan biarkan API crash (sesuai standar deployment) print(f"❌ GAGAL MEMUAT MODEL SECARA LOKAL: {e}") # Anda mungkin perlu menyesuaikan error handling untuk deployment exit(1) # ====================================================== # 1. INISIASI API & MIDDLEWARE # ====================================================== app = FastAPI(title="HF Product Classifier API") # Middleware CORS (Wajib untuk Lovable.dev) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ====================================================== # 2. ENDPOINT PREDIKSI UTAMA (LOGIKA PENUH) # ====================================================== @app.post("/predict") async def predict_product(file: UploadFile = File(...)): """Menerima gambar, menjalankan pipeline, dan mengembalikan JSON.""" try: # 1. Persiapan Gambar contents = await file.read() img_pil = Image.open(BytesIO(contents)).convert("RGB") # 2. Deteksi YOLOv10m results = yolo_model(img_pil, verbose=False) if not results or not results[0].boxes: return JSONResponse(status_code=200, content={"status": "failed", "message": "No object detected."}) # Ambil BBOX terbaik, Crop, Ekstraksi DINOv3 best_box = results[0].boxes.cpu().numpy()[np.argmax(results[0].boxes.cpu().numpy().conf)] x_min, y_min, x_max, y_max = map(int, best_box.xyxy[0]) confidence_score = float(best_box.conf[0]) # Confidence YOLO image_crop = img_pil.crop((x_min, y_min, x_max, y_max)) features = extract_dinov3_features(image_crop) X_pred = features.reshape(1, -1) # 3. Klasifikasi (Probabilitas) Y_proba_list = [clf.predict_proba(X_pred)[0, 1] for clf in all_classifiers] Y_proba = np.array(Y_proba_list) Y_pred_biner = (Y_proba > 0.5).astype(int).reshape(1, -1) class_proba_map = dict(zip(mlb.classes_, Y_proba)) predicted_labels_set = set(mlb.inverse_transform(Y_pred_biner)[0]) final_output = {} # 4. Logika Pemetaan & Fallback for col in label_columns: intersection = predicted_labels_set.intersection(mapping_dict[col]) best_proba = -1 best_label = "UNKNOWN" for label in mapping_dict[col]: proba = class_proba_map.get(label, 0.0) if proba > best_proba: best_proba = proba best_label = label if len(intersection) == 1: final_output[col] = intersection.pop() elif len(intersection) == 0: # UNKNOWN: Fallback ke proba tertinggi jika > 20% if best_proba > 0.20: final_output[col] = f"{best_label} ({best_proba*100:.1f}%)" else: final_output[col] = "UNKNOWN" else: # KONFLIK: Pilih label proba tertinggi final_output[col] = f"CONFLICT -> {best_label} ({best_proba*100:.1f}%)" # 5. Kembalikan Respons JSON Penuh return JSONResponse(status_code=200, content={ "status": "success", "confidence_score": f"{confidence_score*100:.2f}%", "prediction": final_output # Semua 8 atribut terstruktur }) except Exception as e: return JSONResponse(status_code=500, content={"status": "error", "message": str(e)}) # ====================================================== # 3. RUN SERVER (ENTRY POINT UNTUK HUGGING FACE) # ====================================================== # Hugging Face Spaces secara otomatis menjalankan Gunicorn/Uvicorn yang menunjuk ke 'app'. # Di lokal, Anda bisa mengujinya dengan baris ini: if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)