Update app.py

app.py

@@ -98,7 +98,7 @@ class DFUTissueSegNet(nn.Module):
 # تحميل النموذج
 # ======================================================
 def initialize_model():
-    """تحميل DFUTissueSegNet من Google Drive"""
+    """تحميل DFUTissueSegNet من Google Drive مع دعم checkpoint الكامل"""
     global segmenter
     MODEL_URL = "https://drive.google.com/uc?id=1Ovaczsjdp3E-_gYF2pbUibDjPWAC1a6c"
     MODEL_PATH = "best_model_5.pth"
@@ -108,18 +108,34 @@ def initialize_model():
         gdown.download(MODEL_URL, MODEL_PATH, quiet=False)
 
     try:
-        print("🔄 تحميل النموذج...")
+        print("🔄 تحميل DFUTissueSegNet...")
         segmenter = DFUTissueSegNet(num_classes=3)
-        state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
-        segmenter.load_state_dict(state_dict)
+        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+
+        # بعض النماذج تكون محفوظة داخل "state_dict"
+        if "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+        else:
+            state_dict = checkpoint
+
+        # إزالة البادئة "model." أو "module." لو وجدت
+        clean_state = {}
+        for k, v in state_dict.items():
+            nk = k.replace("module.", "").replace("model.", "")
+            clean_state[nk] = v
+
+        segmenter.load_state_dict(clean_state, strict=False)
         segmenter.to(DEVICE)
         segmenter.eval()
         print("✅ تم تحميل DFUTissueSegNet من Google Drive بنجاح.")
     except Exception as e:
-        print(f"❌ خطأ في تحميل النموذج: {e}")
+        print(f"❌ فشل تحميل النموذج: {e}")
+        import traceback
+        traceback.print_exc()
         segmenter = None
 
 
+
 # ======================================================
 # دالة التجزئة
 # ======================================================