Update main.py

main.py

@@ -3,20 +3,16 @@ import numpy as np
 import requests
 import torch
 import firebase_admin
+import os
 from fastapi import FastAPI, BackgroundTasks
 from pydantic import BaseModel
 from ultralytics import YOLO
 from firebase_admin import credentials, firestore
-import os
 
 # --- 0. HUGGING FACE ENVIRONMENT SETUP ---
 os.environ['TORCH_HOME'] = '/tmp/torch_cache'
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/ultralytics_config'
 
-# This is the most reliable way to bypass "Aborted" prompts in 2026
-# It tells Torch Hub to never ask for permission
-torch.hub.help = lambda *args, **kwargs: None 
-
 # --- 1. INITIALIZE MODELS ---
 app = FastAPI()
 
@@ -30,31 +26,48 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 print("🚀 Starting Sahl Express Engine...")
 
 # Load YOLOv8 (2D Segmentation)
-yolo_model = YOLO('best.pt') 
+try:
+    yolo_model = YOLO('best.pt') 
+    print("✅ YOLOv8 Loaded")
+except Exception as e:
+    print(f"❌ YOLO Load Error: {e}")
 
-# Load MiDaS (Depth Estimation) - The "Trust" fix is now built-in
+# --- MIDAS LOADING BLOCK (Bypassing Trusted Repo Prompt) ---
 try:
-    print("📥 Downloading MiDaS models (this may take a minute)...")
-    # We use trust_repo=True here; the global bypass handles the rest
-    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
-    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True)
+    print("📥 Loading MiDaS (Bypassing Trust Prompts)...")
+    
+    # We use skip_validation=True to tell Torch to ignore the repo verification prompt
+    midas = torch.hub.load(
+        "intel-isl/MiDaS", 
+        "MiDaS_small", 
+        trust_repo=True, 
+        skip_validation=True
+    )
     
     midas.to(device)
     midas.eval()
+
+    # Load transforms
+    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
     transform = midas_transforms.small_transform
+    
     print("✅ MiDaS Loaded Successfully")
 except Exception as e:
-    print(f"❌ Error loading MiDaS: {e}")
+    print(f"⚠️ MiDaS Load Failed: {e}")
+    print("🔄 Retrying with fallback method...")
+    # Secondary attempt if the first one fails
+    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", force_reload=True, trust_repo=True)
+    midas.to(device)
+    midas.eval()
 
 # --- 2. FIREBASE SETUP ---
-# CRITICAL: Ensure 'serviceAccount.json' is uploaded to your Space!
 try:
     cred = credentials.Certificate("serviceAccount.json")
     firebase_admin.initialize_app(cred)
     db = firestore.client()
     print("✅ Firebase Connected")
 except Exception as e:
-    print(f"⚠️ Firebase Warning: {e} (App will run but won't save data)")
+    print(f"⚠️ Firebase Error: {e}")
 
 # Tunisian Reference Constants (cm)
 REFERENCE_SIZES = {
@@ -93,7 +106,7 @@ def perform_3d_measurement(image_url: str, delivery_id: str):
         pkg_mask = None
         pkg_w_px, pkg_h_px = None, None
 
-        # 1. Find Ruler (Card/Coin)
+        # 1. Calibration (Find Ruler/ID Card)
         for i, box in enumerate(yolo_results.boxes):
             label = yolo_results.names[int(box.cls[0])]
             if label in REFERENCE_SIZES:
@@ -101,7 +114,7 @@ def perform_3d_measurement(image_url: str, delivery_id: str):
                 pixel_cm_ratio = (x2 - x1) / REFERENCE_SIZES[label]
                 break
 
-        # 2. Find Package
+        # 2. Identification (Find Package)
         for i, box in enumerate(yolo_results.boxes):
             label = yolo_results.names[int(box.cls[0])]
             if label == 'package' and yolo_results.masks is not None:
@@ -111,7 +124,7 @@ def perform_3d_measurement(image_url: str, delivery_id: str):
                 pkg_w_px, pkg_h_px = w, h
                 break
 
-        # 3. Calculate 3D Dimensions
+        # 3. 3D Logic
         if pixel_cm_ratio and pkg_w_px is not None:
             mask_img = np.zeros(depth_map.shape, dtype=np.uint8)
             cv2.fillPoly(mask_img, [pkg_mask.astype(np.int32)], 1)
@@ -121,22 +134,21 @@ def perform_3d_measurement(image_url: str, delivery_id: str):
             dilated = cv2.dilate(mask_img, kernel, iterations=2)
             ground_depth_val = np.median(depth_map[(dilated - mask_img) == 1])
             
-            # Real measurements
             depth_delta = abs(ground_depth_val - pkg_depth_val)
-            real_h = round((depth_delta / pixel_cm_ratio) * 0.5, 1) # Tuning: 0.5
+            real_h = round((depth_delta / pixel_cm_ratio) * 0.5, 1) 
             real_w = round(pkg_w_px / pixel_cm_ratio, 1)
             real_l = round(pkg_h_px / pixel_cm_ratio, 1)
             
             if real_h < 0.5: real_h = 1.0 
             volume = round(real_w * real_l * real_h, 2)
 
-            # 4. Save to Firebase
+            # Update Firebase
             db.collection("orders").document(delivery_id).update({
                 "volume_cm3": volume,
                 "dimensions": f"{real_l}x{real_w}x{real_h} cm",
                 "status": "Measured_3D"
             })
-            print(f"📦 Success: {delivery_id} | {volume} cm3")
+            print(f"📦 Success: {delivery_id} | Vol: {volume}cm3")
 
     except Exception as e:
         print(f"❌ Measurement Error: {e}")
@@ -148,5 +160,4 @@ async def measure_endpoint(request: ImageRequest, background_tasks: BackgroundTa
 
 if __name__ == "__main__":
     import uvicorn
-    # Hugging Face looks for port 7860
     uvicorn.run(app, host="0.0.0.0", port=7860)