Update main.py

main.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 """
-GeeTest4 Solver - Pure FastAPI Version
-100% API endpoint without any UI dependencies
+GeeTest4 Solver - Pure FastAPI Version v1.2.0
+- Updated with optimal thresholds from testing.
+- Fixed YOLOv8 ONNX output transpose bug.
 """
 
 import os
@@ -19,143 +20,111 @@ import numpy as np
 from PIL import Image
 import cv2
 
-# Try to import ONNX (optional)
-ONNX_AVAILABLE = False
+# Coba impor ONNX
 try:
     import onnxruntime as ort
     ONNX_AVAILABLE = True
 except ImportError:
-    pass
+    ONNX_AVAILABLE = False
 
-# --- Configuration ---
+# ===================================================================
+# KONFIGURASI
+# ===================================================================
 MODEL_PATH = "best_model.onnx"
-YAML_PATH = "data.yaml" 
-CONFIDENCE_THRESHOLD = 0.3
-NMS_IOU_THRESHOLD = 0.5
+YAML_PATH = "data.yaml"
 API_KEY = os.getenv("GEETEST4_API_KEY", "ADMINCKV005")
 
-# Global variables
+# MODIFIKASI: Menggunakan nilai optimal dari hasil tes interaktif
+CONFIDENCE_THRESHOLD = 0.25
+NMS_IOU_THRESHOLD = 0.0
+
+# Variabel Global
 model_session = None
 CLASS_NAMES = []
 
-# --- Setup Logging ---
-logging.basicConfig(level=logging.INFO)
+# Setup Logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# --- Pydantic Models ---
+# Pydantic Models
 class PredictRequest(BaseModel):
     data: List[str]
-
-# MODIFIKASI: Tipe data untuk bounding box
 BoundingBox = Dict[str, int]
 
 def verify_api_key(api_key: str) -> bool:
-    """Verify API key"""
     return api_key == API_KEY
 
-def smart_cv_model(image_np: np.ndarray) -> Tuple[int, float, Union[BoundingBox, None]]:
-    """Computer Vision model - sekarang mengembalikan bounding box"""
-    height, width = image_np.shape[:2]
-    try:
-        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-        _, thresh = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
-        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        if contours:
-            largest_contour = max(contours, key=cv2.contourArea)
-            x, y, w, h = cv2.boundingRect(largest_contour)
-            center_x = x + w // 2
-            
-            # MODIFIKASI: Siapkan data bounding box untuk dikembalikan
-            bbox = {'x': x, 'y': y, 'w': w, 'h': h}
-            
-            center_x = max(int(width * 0.1), min(center_x, int(width * 0.9)))
-            area_ratio = cv2.contourArea(largest_contour) / (width * height)
-            confidence = min(0.9, max(0.6, area_ratio * 10))
-            
-            logger.info(f"CV Model: target at x={center_x}, confidence={confidence:.3f}")
-            # MODIFIKASI: Kembalikan bbox
-            return center_x, confidence, bbox
-        else:
-            # Fallback jika tidak ada kontur
-            random.seed(hash(image_np.tobytes()) % 2**31)
-            target_x = int(width * (0.45 + random.random() * 0.3))
-            confidence = 0.65 + random.random() * 0.15
-            logger.info(f"CV Model (rule-based): target at x={target_x}, confidence={confidence:.3f}")
-            return target_x, confidence, None
-            
-    except Exception as e:
-        logger.warning(f"CV processing failed, using safe fallback: {e}")
-        center_x = int(width * 0.6)
-        return center_x, 0.7, None
+def preprocess_for_onnx(image: np.ndarray, input_size: int = 640):
+    height, width, _ = image.shape
+    r = min(input_size / width, input_size / height)
+    new_width, new_height = int(width * r), int(height * r)
+    resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
+    d_w, d_h = (input_size - new_width) // 2, (input_size - new_height) // 2
+    padded_image = np.full((input_size, input_size, 3), 114, dtype=np.uint8)
+    padded_image[d_h:new_height + d_h, d_w:new_width + d_w, :] = resized_image
+    input_tensor = (padded_image.astype(np.float32) / 255.0).transpose(2, 0, 1)
+    return np.expand_dims(input_tensor, axis=0), r, d_w, d_h
 
 def process_image_onnx(image_np: np.ndarray) -> Tuple[int, float, Union[BoundingBox, None]]:
-    """Process image with ONNX model - sekarang mengembalikan bounding box"""
-    height, width = image_np.shape[:2]
+    """Memproses gambar dengan model ONNX."""
     try:
-        max_size = 640
-        ratio = min(max_size / width, max_size / height)
-        new_width, new_height = int(width * ratio), int(height * ratio)
-        resized = cv2.resize(image_np, (new_width, new_height))
-        
-        dw, dh = (max_size - new_width) // 2, (max_size - new_height) // 2
-        padded = cv2.copyMakeBorder(resized, dh, max_size - new_height - dh, dw, max_size - new_width - dw, cv2.BORDER_CONSTANT, value=(114, 114, 114))
-        
-        input_tensor = (padded.astype(np.float32) / 255.0).transpose(2, 0, 1)
-        input_tensor = np.expand_dims(input_tensor, axis=0)
+        input_tensor, ratio, dw, dh = preprocess_for_onnx(image_np)
         
+        # Jalankan inferensi
         outputs = model_session.run(None, {model_session.get_inputs()[0].name: input_tensor})
-        preds = outputs[0][0]
         
-        if len(preds) == 0:
-            return 0, 0.0, None
+        # ==================================================================
+        # PERBAIKAN KRUSIAL: Tambahkan .T untuk transpose output model YOLOv8
+        # ==================================================================
+        raw_predictions = outputs[0][0].T
+
+        # Ambil skor confidence (sekarang berada di kolom yang benar)
+        scores = raw_predictions[:, 4]
+        valid_indices = scores > CONFIDENCE_THRESHOLD
         
-        box_scores = preds[:, 4:]
-        max_scores = np.max(box_scores, axis=1)
+        boxes_raw = raw_predictions[valid_indices, :4]
+        scores = scores[valid_indices]
         
-        valid_preds = max_scores > CONFIDENCE_THRESHOLD
-        if not np.any(valid_preds):
+        if len(boxes_raw) == 0:
             return 0, 0.0, None
         
-        preds = preds[valid_preds]
-        max_scores = max_scores[valid_preds]
-        
-        boxes_raw = preds[:, :4]
-        boxes_raw[:, 0] = (boxes_raw[:, 0] - dw) / ratio
-        boxes_raw[:, 1] = (boxes_raw[:, 1] - dh) / ratio
-        boxes_raw[:, 2] /= ratio
-        boxes_raw[:, 3] /= ratio
-        
-        x1, y1 = boxes_raw[:, 0] - boxes_raw[:, 2] / 2, boxes_raw[:, 1] - boxes_raw[:, 3] / 2
-        x2, y2 = boxes_raw[:, 0] + boxes_raw[:, 2] / 2, boxes_raw[:, 1] + boxes_raw[:, 3] / 2
-        
-        boxes_processed = np.column_stack((x1, y1, x2, y2)).astype(np.float32)
-        indices = cv2.dnn.NMSBoxes(boxes_processed, max_scores, CONFIDENCE_THRESHOLD, NMS_IOU_THRESHOLD)
+        # Konversi box dari (center_x, center_y, w, h) ke (x1, y1, x2, y2)
+        x1 = boxes_raw[:, 0] - boxes_raw[:, 2] / 2
+        y1 = boxes_raw[:, 1] - boxes_raw[:, 3] / 2
+        x2 = boxes_raw[:, 0] + boxes_raw[:, 2] / 2
+        y2 = boxes_raw[:, 1] + boxes_raw[:, 3] / 2
+        boxes_for_nms = np.column_stack((x1, y1, x2, y2)).astype(np.float32)
+
+        # Terapkan Non-Max Suppression
+        indices = cv2.dnn.NMSBoxes(boxes_for_nms, scores, CONFIDENCE_THRESHOLD, NMS_IOU_THRESHOLD)
         
         if len(indices) == 0:
             return 0, 0.0, None
-        
-        best_idx = indices.flatten()[0]
-        best_box = boxes_processed[best_idx]
-        best_score = max_scores[best_idx]
-        center_x = int((best_box[0] + best_box[2]) / 2)
 
-        # MODIFIKASI: Konversi dari x1,y1,x2,y2 ke x,y,w,h
-        x = int(best_box[0])
-        y = int(best_box[1])
-        w = int(best_box[2] - best_box[0])
-        h = int(best_box[3] - best_box[1])
-        bbox = {'x': x, 'y': y, 'w': w, 'h': h}
+        # Ambil deteksi terbaik (dengan skor tertinggi setelah NMS)
+        indices = indices.flatten()
+        best_idx = indices[np.argmax(scores[indices])]
+        
+        best_box_coords = boxes_for_nms[best_idx]
+        best_score = scores[best_idx]
+        
+        # Konversi koordinat kembali ke ukuran gambar asli
+        x1_orig = int((best_box_coords[0] - dw) / ratio)
+        y1_orig = int((best_box_coords[1] - dh) / ratio)
+        x2_orig = int((best_box_coords[2] - dw) / ratio)
+        y2_orig = int((best_box_coords[3] - dh) / ratio)
+        
+        center_x = (x1_orig + x2_orig) // 2
+        bbox = {'x': x1_orig, 'y': y1_orig, 'w': x2_orig - x1_orig, 'h': y2_orig - y1_orig}
         
-        # MODIFIKASI: Kembalikan bbox
         return center_x, float(best_score), bbox
         
     except Exception as e:
-        logger.error(f"Error in ONNX processing: {e}")
+        logger.error(f"Error dalam pemrosesan ONNX: {e}")
         return 0, 0.0, None
 
 def load_model():
-    """Load ONNX model and class names"""
+    """Memuat model ONNX."""
     global model_session, CLASS_NAMES
     try:
         if os.path.exists(YAML_PATH):
@@ -163,68 +132,56 @@ def load_model():
                 CLASS_NAMES = yaml.safe_load(f).get('names', ['Target'])
         else:
             CLASS_NAMES = ['Target']
-        logger.info(f"Loaded {len(CLASS_NAMES)} classes: {CLASS_NAMES}")
         
         if ONNX_AVAILABLE and os.path.exists(MODEL_PATH):
             model_session = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
-            logger.info("✅ ONNX model loaded successfully")
+            logger.info("✅ Model ONNX berhasil dimuat.")
         else:
-            logger.warning("ONNX model not found or ONNX runtime not available. Using CV model.")
             model_session = None
+            logger.critical("❌ GAGAL: Model ONNX tidak ditemukan atau onnxruntime tidak terinstal.")
+            
     except Exception as e:
-        logger.error(f"Failed to load model: {e}")
+        logger.error(f"FATAL: Gagal memuat model: {e}")
         model_session = None
 
 def base64_to_numpy(base64_string: str) -> np.ndarray:
-    """Convert base64 to numpy array"""
     try:
         if base64_string.startswith('data:image'):
             base64_string = base64_string.split(',')[1]
         image_data = base64.b64decode(base64_string)
         return np.array(Image.open(io.BytesIO(image_data)).convert('RGB'))
     except Exception as e:
-        logger.error(f"Error converting base64: {e}")
-        raise ValueError("Invalid image data")
+        logger.error(f"Error saat konversi base64: {e}")
+        raise ValueError("Data gambar tidak valid")
 
 def solve_geetest4_api(background_image: str, api_key: str):
-    """Pure API endpoint function - sekarang mengembalikan bounding box"""
+    """Fungsi endpoint API utama."""
     try:
         if not verify_api_key(api_key):
-            # MODIFIKASI: Tambah None untuk konsistensi format
-            return ["❌ Invalid API key", 0, 0.0, None]
+            return ["❌ Kunci API tidak valid", 0, 0.0, None]
         
         image_np = base64_to_numpy(background_image)
-        bbox = None # Default value
-
-        if model_session:
-            # MODIFIKASI: Tangkap bbox dari return value
+        
+        if model_session is not None:
             target_x, confidence, bbox = process_image_onnx(image_np)
-            if confidence < CONFIDENCE_THRESHOLD:
-                logger.info("ONNX confidence too low, using CV fallback")
-                target_x, confidence, bbox = smart_cv_model(image_np)
-                model_type = "CV"
-            else:
-                model_type = "ONNX"
+            model_type = "ONNX"
         else:
-            # MODIFIKASI: Tangkap bbox dari return value
-            target_x, confidence, bbox = smart_cv_model(image_np)
-            model_type = "CV"
+            return ["❌ Model tidak dimuat", 0, 0.0, None]
         
-        if target_x > 0 and confidence >= 0.5:
-            # MODIFIKASI: Sertakan bbox dalam respons sukses
-            return [f"✅ Success! Target at x={target_x} (Model: {model_type})", target_x, confidence, bbox]
+        if target_x > 0 and bbox is not None:
+            return [f"✅ Sukses! Target di x={target_x} (Model: {model_type})", target_x, confidence, bbox]
         else:
-            fallback_x = int(image_np.shape[1] * 0.6)
-            # MODIFIKASI: Tambah None untuk konsistensi format
-            return [f"✅ Fallback position x={fallback_x}", fallback_x, 0.7, None]
+            return [f"⚠️ Tidak ada target terdeteksi dengan threshold saat ini.", 0, 0.0, None]
         
     except Exception as e:
-        logger.error(f"API Error: {e}")
-        # MODIFIKASI: Tambah None untuk konsistensi format
-        return [f"⚠️ Error, using fallback position", 200, 0.6, None]
+        logger.error(f"Error API: {e}")
+        return [f"⚠️ Error server, menggunakan posisi fallback", 200, 0.6, None]
 
+# Inisialisasi model saat startup
 load_model()
-app = FastAPI(title="GeeTest4 Solver API", version="1.1.0", docs_url=None, redoc_url=None)
+
+# --- Aplikasi FastAPI ---
+app = FastAPI(title="GeeTest4 Solver API", version="1.2.0", docs_url=None, redoc_url=None)
 
 @app.get("/")
 async def root():
@@ -232,22 +189,29 @@ async def root():
 
 @app.post("/api/predict")
 async def predict(request: PredictRequest):
-    """Main API endpoint for GeeTest4 solving"""
+    """Endpoint utama untuk prediksi."""
     try:
         if len(request.data) < 2:
-            raise HTTPException(status_code=400, detail="Invalid request format")
+            raise HTTPException(status_code=400, detail="Format request tidak valid")
         
         background_image, api_key = request.data[0], request.data[1]
         result = solve_geetest4_api(background_image, api_key)
         return {"data": result}
+        
     except Exception as e:
         logger.error(f"API Error: {e}")
-        return JSONResponse(status_code=500, content={"data": ["❌ Server error", 0, 0.0, None]})
+        return JSONResponse(status_code=500, content={"data": ["❌ Error internal server", 0, 0.0, None]})
 
 @app.get("/health")
 async def health_check():
     return {"status": "healthy", "model_loaded": model_session is not None}
 
+# Menjalankan aplikasi
 if __name__ == "__main__":
-    logger.info("🚀 Starting GeeTest4 Pure FastAPI Server")
-    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)), log_level="info")
+    logger.info("🚀 Memulai Server FastAPI GeeTest4...")
+    uvicorn.run(
+        "__main__:app",
+        host="0.0.0.0",
+        port=int(os.getenv("PORT", 7860)),
+        log_level="info"
+    )