Update app.py

app.py

@@ -52,6 +52,10 @@ for _mod in ["OpenGL", "OpenGL.GL", "OpenGL.GL.framebufferobjects",
 
 os.environ["PYOPENGL_PLATFORM"] = "osmesa"
 
+# --- Router Model Classes ---
+CLASSES = {0: "letter", 1: "number"}
+IMG_SIZE = 64
+
 # --- Hugging Face Model Integration ---
 REPO_ID = "SondosM/api_GP"
 
@@ -76,14 +80,11 @@ get_hf_file("mano_data/mano_data/mano_mean_params.npz", is_mano=True)
 get_hf_file("mano_data/mano_data/MANO_LEFT.pkl",        is_mano=True)
 get_hf_file("mano_data/mano_data/MANO_RIGHT.pkl",       is_mano=True)
 
-WILOR_REPO_PATH = "./WiLoR"
-WILOR_CKPT      = get_hf_file("pretrained_models/pretrained_models/wilor_final.ckpt")
-WILOR_CFG       = get_hf_file("pretrained_models/pretrained_models/model_config.yaml")
-DETECTOR_PATH   = get_hf_file("pretrained_models/pretrained_models/detector.pt")
-
-# ─── الفرق الأساسي: الكود الأول كان بيحمّل classifier.pkl من مسار محلي ثابت
-# بدل ما يحمّله من HF زي باقي الملفات ─────────────────────────────────────────
-CLASSIFIER_PATH   = get_hf_file("classifier.pkl")
+WILOR_REPO_PATH   = "./WiLoR"
+WILOR_CKPT        = get_hf_file("pretrained_models/pretrained_models/wilor_final.ckpt")
+WILOR_CFG         = get_hf_file("pretrained_models/pretrained_models/model_config.yaml")
+DETECTOR_PATH     = get_hf_file("pretrained_models/pretrained_models/detector.pt")
+ROUTER_MODEL_PATH = get_hf_file("router_model.keras")
 MLP_LETTERS_PATH  = get_hf_file("MLP_letters.pkl")
 MLP_NUMBERS_PATH  = get_hf_file("MLP_numbers.pkl")
 
@@ -94,19 +95,20 @@ WILOR_TRANSFORM = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-wilor_model   = None
-yolo_detector = None
-classifier    = None
-mlp_letters   = None
-mlp_numbers   = None
+wilor_model        = None
+yolo_detector      = None
+router_model_keras = None
+mlp_letters        = None
+mlp_numbers        = None
 
 
 def load_models():
-    global wilor_model, yolo_detector, classifier, mlp_letters, mlp_numbers
+    global wilor_model, yolo_detector, router_model_keras, mlp_letters, mlp_numbers
 
     sys.path.insert(0, WILOR_REPO_PATH)
     from wilor.models import load_wilor
     from ultralytics import YOLO
+    from tensorflow.keras.models import load_model
 
     print(f"Loading WiLoR on {DEVICE}...")
     wilor_model, _ = load_wilor(checkpoint_path=WILOR_CKPT, cfg_path=WILOR_CFG)
@@ -116,8 +118,10 @@ def load_models():
     print("Loading YOLO detector...")
     yolo_detector = YOLO(DETECTOR_PATH)
 
-    print("Loading classifiers...")
-    classifier  = joblib.load(CLASSIFIER_PATH)
+    print("Loading router model (Keras)...")
+    router_model_keras = load_model(ROUTER_MODEL_PATH)
+
+    print("Loading MLP classifiers...")
     mlp_letters = joblib.load(MLP_LETTERS_PATH)
     mlp_numbers = joblib.load(MLP_NUMBERS_PATH)
 
@@ -189,8 +193,6 @@ def read_image_from_upload(file_bytes: bytes) -> np.ndarray:
 
 
 def _align_features(model, features: np.ndarray) -> np.ndarray:
-    # بعض الـ models (Pipeline مع StandardScaler) مش بيحفظوا feature_names_in_
-    # فبنستخدم n_features_in_ بدلها ونرجع array عادية مش DataFrame
     if hasattr(model, "feature_names_in_"):
         expected_cols = model.feature_names_in_
         vec           = np.zeros(len(expected_cols))
@@ -198,33 +200,25 @@ def _align_features(model, features: np.ndarray) -> np.ndarray:
         vec[:limit]   = features[:limit]
         return pd.DataFrame([vec], columns=expected_cols)
     else:
-        n = model.n_features_in_
-        vec = np.zeros(n)
+        n     = model.n_features_in_
+        vec   = np.zeros(n)
         limit = min(len(features), n)
         vec[:limit] = features[:limit]
         return vec.reshape(1, -1)
 
 
-def run_two_stage(features: np.ndarray) -> dict:
-    # Stage 1: letter or number?
-    feat_df  = _align_features(classifier, features)
-    category = str(classifier.predict(feat_df)[0])
-    cat_conf = float(classifier.predict_proba(feat_df)[0].max())
+def run_two_stage(features: np.ndarray, crop_rgb: np.ndarray) -> dict:
+    # Stage 1: router_model.keras يحدد حرف (0) أو رقم (1)
+    img_resized = cv2.resize(crop_rgb, (IMG_SIZE, IMG_SIZE))
+    img_array   = np.expand_dims(img_resized, axis=0).astype("float32") / 255.0
 
-    # Stage 2: which sign exactly?
-    cat = category.lower().strip()
-    if cat in ("letter", "letters", "حرف", "حروف"):
-        model = mlp_letters
-    elif cat in ("number", "numbers", "digit", "digits", "رقم", "أرقام", "ارقام"):
-        model = mlp_numbers
-    else:
-        # fallback: pick whichever is more confident
-        feat_l  = _align_features(mlp_letters, features)
-        feat_n  = _align_features(mlp_numbers, features)
-        prob_l  = float(mlp_letters.predict_proba(feat_l)[0].max())
-        prob_n  = float(mlp_numbers.predict_proba(feat_n)[0].max())
-        model   = mlp_letters if prob_l >= prob_n else mlp_numbers
+    prob    = float(router_model_keras.predict(img_array, verbose=0)[0][0])
+    cls_idx = 1 if prob >= 0.5 else 0
+    category = CLASSES[cls_idx]
+    cat_conf = prob if cls_idx == 1 else 1.0 - prob
 
+    # Stage 2: اختار الموديل الصح بناءً على النتيجة
+    model   = mlp_letters if category == "letter" else mlp_numbers
     feat_df = _align_features(model, features)
     label   = str(model.predict(feat_df)[0])
     conf    = float(model.predict_proba(feat_df)[0].max())
@@ -256,8 +250,8 @@ async def predict(file: UploadFile = File(...)):
     if not results[0].boxes:
         raise HTTPException(status_code=422, detail="No hand detected.")
 
-    box      = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
-    label_id = int(results[0].boxes.cls[0].cpu().item())
+    box       = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
+    label_id  = int(results[0].boxes.cls[0].cpu().item())
     hand_side = "left" if label_id == 0 else "right"
 
     h, w = img_rgb.shape[:2]
@@ -271,7 +265,7 @@ async def predict(file: UploadFile = File(...)):
     if features is None:
         raise HTTPException(status_code=500, detail="Feature extraction failed.")
 
-    result = run_two_stage(features)
+    result = run_two_stage(features, crop)
     return JSONResponse({**result, "hand_side": hand_side, "bbox": [int(x1), int(y1), int(x2), int(y2)]})
 
 
@@ -285,8 +279,8 @@ async def predict_with_skeleton(file: UploadFile = File(...)):
     if not results[0].boxes:
         raise HTTPException(status_code=422, detail="No hand detected.")
 
-    box      = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
-    label_id = int(results[0].boxes.cls[0].cpu().item())
+    box       = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
+    label_id  = int(results[0].boxes.cls[0].cpu().item())
     hand_side = "left" if label_id == 0 else "right"
 
     h, w = img_rgb.shape[:2]
@@ -296,7 +290,7 @@ async def predict_with_skeleton(file: UploadFile = File(...)):
     features = extract_features(crop)
     joints   = get_3d_joints(crop)
 
-    result   = run_two_stage(features)
+    result   = run_two_stage(features, crop)
     _, buf   = cv2.imencode(".png", cv2.cvtColor(crop, cv2.COLOR_RGB2BGR))
     crop_b64 = base64.b64encode(buf).decode("utf-8")