Update main.py

main.py

@@ -10,7 +10,7 @@ import io
 import logging
 import random
 import yaml
-from typing import Tuple, List
+from typing import Tuple, List, Dict, Union
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
@@ -46,176 +46,131 @@ logger = logging.getLogger(__name__)
 class PredictRequest(BaseModel):
     data: List[str]
 
-class PredictResponse(BaseModel):
-    data: List
+# 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]:
-    """Computer Vision model - ALWAYS succeeds"""
+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:
-        height, width = image_np.shape[:2]
-        
-        # Convert to grayscale for analysis
         gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-        
-        # Apply threshold to find dark areas (potential missing pieces)
         _, thresh = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
-        
-        # Find contours
         contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
         if contours:
-            # Find the largest contour (likely the missing piece)
             largest_contour = max(contours, key=cv2.contourArea)
-            
-            # Get bounding box
             x, y, w, h = cv2.boundingRect(largest_contour)
-            
-            # Calculate center X
             center_x = x + w // 2
             
-            # Ensure position is within valid range
-            center_x = max(int(width * 0.1), min(center_x, int(width * 0.9)))
+            # MODIFIKASI: Siapkan data bounding box untuk dikembalikan
+            bbox = {'x': x, 'y': y, 'w': w, 'h': h}
             
-            # Calculate confidence
+            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}")
-            return center_x, confidence
+            # MODIFIKASI: Kembalikan bbox
+            return center_x, confidence, bbox
         else:
-            # Fallback: deterministic positioning
+            # 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
+            return target_x, confidence, None
             
     except Exception as e:
         logger.warning(f"CV processing failed, using safe fallback: {e}")
-        # Safe fallback - ALWAYS works
         center_x = int(width * 0.6)
-        return center_x, 0.7
+        return center_x, 0.7, None
 
-def process_image_onnx(image_np: np.ndarray) -> Tuple[int, float]:
-    """Process image with ONNX model"""
+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]
     try:
-        # Preprocess image for ONNX
-        height, width = image_np.shape[:2]
-        
-        # Resize and pad to 640x640
         max_size = 640
         ratio = min(max_size / width, max_size / height)
-        new_width = int(width * ratio)
-        new_height = int(height * ratio)
-        
+        new_width, new_height = int(width * ratio), int(height * ratio)
         resized = cv2.resize(image_np, (new_width, new_height))
         
-        # Pad to square
-        dw = (max_size - new_width) // 2
-        dh = (max_size - new_height) // 2
+        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))
         
-        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)
         
-        # Normalize and prepare for ONNX
-        input_tensor = padded.astype(np.float32) / 255.0
-        input_tensor = np.transpose(input_tensor, (2, 0, 1))  # HWC to CHW
-        input_tensor = np.expand_dims(input_tensor, axis=0)   # Add batch dimension
-        
-        # Run inference
         outputs = model_session.run(None, {model_session.get_inputs()[0].name: input_tensor})
-        preds = outputs[0][0]  # Remove batch dimension
+        preds = outputs[0][0]
         
-        # Post-process predictions
         if len(preds) == 0:
-            return 0, 0.0
+            return 0, 0.0, None
         
-        # Extract boxes and scores
         box_scores = preds[:, 4:]
         max_scores = np.max(box_scores, axis=1)
         
-        # Filter by confidence
         valid_preds = max_scores > CONFIDENCE_THRESHOLD
         if not np.any(valid_preds):
-            return 0, 0.0
+            return 0, 0.0, None
         
         preds = preds[valid_preds]
         max_scores = max_scores[valid_preds]
         
-        # Process bounding boxes
         boxes_raw = preds[:, :4]
-        
-        # Scale back to original image
         boxes_raw[:, 0] = (boxes_raw[:, 0] - dw) / ratio
         boxes_raw[:, 1] = (boxes_raw[:, 1] - dh) / ratio
         boxes_raw[:, 2] /= ratio
         boxes_raw[:, 3] /= ratio
         
-        # Convert to x1, y1, x2, y2 for NMS
-        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
+        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)
-        
-        # Apply NMS
         indices = cv2.dnn.NMSBoxes(boxes_processed, max_scores, CONFIDENCE_THRESHOLD, NMS_IOU_THRESHOLD)
         
         if len(indices) == 0:
-            return 0, 0.0
+            return 0, 0.0, None
         
-        # Get best detection
         best_idx = indices.flatten()[0]
         best_box = boxes_processed[best_idx]
         best_score = max_scores[best_idx]
-        
-        # Calculate center x position
         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}
         
-        return center_x, float(best_score)
+        # MODIFIKASI: Kembalikan bbox
+        return center_x, float(best_score), bbox
         
     except Exception as e:
         logger.error(f"Error in ONNX processing: {e}")
-        return 0, 0.0
+        return 0, 0.0, None
 
 def load_model():
     """Load ONNX model and class names"""
     global model_session, CLASS_NAMES
-    
     try:
-        # Load class names
         if os.path.exists(YAML_PATH):
             with open(YAML_PATH, "r", encoding="utf-8") as f:
-                data = yaml.safe_load(f)
-                CLASS_NAMES = data.get('names', ['Target'])
+                CLASS_NAMES = yaml.safe_load(f).get('names', ['Target'])
         else:
             CLASS_NAMES = ['Target']
-        
         logger.info(f"Loaded {len(CLASS_NAMES)} classes: {CLASS_NAMES}")
         
-        # Load ONNX model
         if ONNX_AVAILABLE and os.path.exists(MODEL_PATH):
-            try:
-                providers = ['CPUExecutionProvider']
-                model_session = ort.InferenceSession(MODEL_PATH, providers=providers)
-                logger.info("✅ ONNX model loaded successfully")
-            except Exception as e:
-                logger.error(f"Failed to load ONNX model: {e}")
-                logger.info("Will use CV model instead")
-                model_session = None
+            model_session = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
+            logger.info("✅ ONNX model loaded successfully")
         else:
-            if not ONNX_AVAILABLE:
-                logger.warning("ONNX Runtime not available, using CV model")
-            elif not os.path.exists(MODEL_PATH):
-                logger.warning(f"Model file {MODEL_PATH} not found, using CV model")
+            logger.warning("ONNX model not found or ONNX runtime not available. Using CV model.")
             model_session = None
-            
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
         model_session = None
@@ -225,107 +180,74 @@ def base64_to_numpy(base64_string: str) -> np.ndarray:
     try:
         if base64_string.startswith('data:image'):
             base64_string = base64_string.split(',')[1]
-        
         image_data = base64.b64decode(base64_string)
-        image = Image.open(io.BytesIO(image_data))
-        return np.array(image.convert('RGB'))
+        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")
 
 def solve_geetest4_api(background_image: str, api_key: str):
-    """Pure API endpoint function - always returns consistent format"""
+    """Pure API endpoint function - sekarang mengembalikan bounding box"""
     try:
-        # Verify API key
         if not verify_api_key(api_key):
-            return ["❌ Invalid API key", 0, 0.0]
+            # MODIFIKASI: Tambah None untuk konsistensi format
+            return ["❌ Invalid API key", 0, 0.0, None]
         
-        # Convert image
         image_np = base64_to_numpy(background_image)
-        
-        # Process with appropriate model
-        if model_session is not None:
-            target_x, confidence = process_image_onnx(image_np)
-            
-            # If ONNX fails, fallback to CV
+        bbox = None # Default value
+
+        if model_session:
+            # MODIFIKASI: Tangkap bbox dari return value
+            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 = smart_cv_model(image_np)
+                target_x, confidence, bbox = smart_cv_model(image_np)
                 model_type = "CV"
             else:
                 model_type = "ONNX"
         else:
-            # Use CV model directly
-            target_x, confidence = smart_cv_model(image_np)
+            # MODIFIKASI: Tangkap bbox dari return value
+            target_x, confidence, bbox = smart_cv_model(image_np)
             model_type = "CV"
         
-        # Always succeed with minimum confidence
         if target_x > 0 and confidence >= 0.5:
-            return [f"✅ Success! Target at x={target_x} (Model: {model_type})", target_x, confidence]
+            # MODIFIKASI: Sertakan bbox dalam respons sukses
+            return [f"✅ Success! Target at x={target_x} (Model: {model_type})", target_x, confidence, bbox]
         else:
-            # Fallback for edge cases
             fallback_x = int(image_np.shape[1] * 0.6)
-            return [f"✅ Fallback position x={fallback_x}", fallback_x, 0.7]
+            # MODIFIKASI: Tambah None untuk konsistensi format
+            return [f"✅ Fallback position x={fallback_x}", fallback_x, 0.7, None]
         
     except Exception as e:
         logger.error(f"API Error: {e}")
-        # Even on error, return a usable position
-        return [f"⚠️ Error, using fallback position", 200, 0.6]
+        # MODIFIKASI: Tambah None untuk konsistensi format
+        return [f"⚠️ Error, using fallback position", 200, 0.6, None]
 
-# Initialize model on startup
 load_model()
-
-# --- FastAPI App ---
-app = FastAPI(
-    title="GeeTest4 Solver API",
-    description="Pure API endpoint for GeeTest4 captcha solving",
-    version="1.0.0",
-    docs_url=None,  # Disable docs
-    redoc_url=None  # Disable redoc
-)
+app = FastAPI(title="GeeTest4 Solver API", version="1.1.0", docs_url=None, redoc_url=None)
 
 @app.get("/")
 async def root():
-    """Root endpoint - returns 404 to hide the service"""
     raise HTTPException(status_code=404, detail="Not Found")
 
 @app.post("/api/predict")
 async def predict(request: PredictRequest):
     """Main API endpoint for GeeTest4 solving"""
     try:
-        # Validate request
         if len(request.data) < 2:
             raise HTTPException(status_code=400, detail="Invalid request format")
         
-        background_image = request.data[0]
-        api_key = request.data[1]
-        
-        # Process request
+        background_image, api_key = request.data[0], request.data[1]
         result = solve_geetest4_api(background_image, api_key)
-        
-        # Return response in Gradio format for compatibility
         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]}
-        )
+        return JSONResponse(status_code=500, content={"data": ["❌ Server error", 0, 0.0, None]})
 
 @app.get("/health")
 async def health_check():
-    """Health check endpoint"""
     return {"status": "healthy", "model_loaded": model_session is not None}
 
-# Launch app
 if __name__ == "__main__":
     logger.info("🚀 Starting GeeTest4 Pure FastAPI Server")
-    
-    # Launch with uvicorn
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=int(os.getenv("PORT", 7860)),
-        log_level="info"
-    )
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)), log_level="info")