Update app.py

app.py

@@ -1,10 +1,7 @@
 # -*- coding: utf-8 -*-
-# -*- coding: utf-8 -*-
 """
-واجهة موحّدة لتحليل صورة القدم باستخدام DFUTissueSegNet (من Google Drive)
-- ألوان متعددة: قرحة=أحمر، Slough=أصفر، نخر=أسود
-- حساب نسب كل نوع + مستوى خطورة إجمالي
-- Legend داخل Gradio
+تحليل قرحة القدم باستخدام Unet + EfficientNet-b0
+النموذج من Google Drive (best_model_5.pth)
 """
 
 import os
@@ -12,233 +9,165 @@ import cv2
 import gdown
 import numpy as np
 from PIL import Image
-
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
 import gradio as gr
+import segmentation_models_pytorch as smp
 
-# ================================
-# إعدادات عامة
-# ================================
+# =========================================================
+# الإعدادات العامة
+# =========================================================
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-IMG_SIZE = 512                    # حجم الإدخال كما في تدريب DFUTissueSegNet
-THRESHOLD = 0.4                   # عتبة تحويل الاحتمال إلى قناع (جرّبي 0.35..0.6 حسب بياناتك)
+IMG_SIZE = 512
+THRESHOLD = 0.35
 MODEL_PATH = "best_model_5.pth"
-
-# غيّري الـ File ID أدناه لملفّك على Google Drive عند الحاجة
-# مثال: https://drive.google.com/file/d/FILE_ID/view => استخدمي: https://drive.google.com/uc?id=FILE_ID
 MODEL_URL = "https://drive.google.com/uc?id=1Ovaczsjdp3E-_gYF2pbUibDjPWAC1a6c"
 
-CLASS_NAMES = ["قرحة (Granulation)", "Slough (أنسجة ميتة جزئيًا)", "نخر (Necrotic)"]
+CLASS_NAMES = ["قرحة (Granulation)", "Slough", "نخر (Necrosis)"]
 CLASS_COLORS = {
-    "قرحة (Granulation)": (255, 0, 0),     # أحمر
-    "Slough (أنسجة ميتة جزئيًا)": (255, 255, 0),  # أصفر
-    "نخر (Necrotic)": (0, 0, 0)            # أسود
+    "قرحة (Granulation)": (255, 0, 0),   # أحمر
+    "Slough": (255, 255, 0),             # أصفر
+    "نخر (Necrosis)": (0, 0, 0)          # أسود
 }
 
-# ================================
-# نموذج DFUTissueSegNet
-# ================================
-class ConvBlock(nn.Module):
-    def __init__(self, in_ch, out_ch):
-        super().__init__()
-        self.block = nn.Sequential(
-            nn.Conv2d(in_ch, out_ch, 3, padding=1),
-            nn.BatchNorm2d(out_ch),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(out_ch, out_ch, 3, padding=1),
-            nn.BatchNorm2d(out_ch),
-            nn.ReLU(inplace=True),
-        )
-    def forward(self, x):
-        return self.block(x)
-
-class DFUTissueSegNet(nn.Module):
-    def __init__(self, num_classes=3):
-        super().__init__()
-        self.enc1 = ConvBlock(3, 64);   self.pool1 = nn.MaxPool2d(2)
-        self.enc2 = ConvBlock(64, 128); self.pool2 = nn.MaxPool2d(2)
-        self.enc3 = ConvBlock(128, 256);self.pool3 = nn.MaxPool2d(2)
-        self.enc4 = ConvBlock(256, 512);self.pool4 = nn.MaxPool2d(2)
-
-        self.center = ConvBlock(512, 1024)
-
-        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
-        self.dec4 = ConvBlock(1024, 512)
-        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
-        self.dec3 = ConvBlock(512, 256)
-        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
-        self.dec2 = ConvBlock(256, 128)
-        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
-        self.dec1 = ConvBlock(128, 64)
-
-        self.final = nn.Conv2d(64, num_classes, 1)
-
-    def forward(self, x):
-        e1 = self.enc1(x)
-        e2 = self.enc2(self.pool1(e1))
-        e3 = self.enc3(self.pool2(e2))
-        e4 = self.enc4(self.pool3(e3))
-        c  = self.center(self.pool4(e4))
-
-        d4 = self.up4(c); d4 = torch.cat([d4, e4], dim=1); d4 = self.dec4(d4)
-        d3 = self.up3(d4); d3 = torch.cat([d3, e3], dim=1); d3 = self.dec3(d3)
-        d2 = self.up2(d3); d2 = torch.cat([d2, e2], dim=1); d2 = self.dec2(d2)
-        d1 = self.up1(d2); d1 = torch.cat([d1, e1], dim=1); d1 = self.dec1(d1)
-
-        # مخرجات احتمالية 0..1 لكل قناة
-        return torch.sigmoid(self.final(d1))
-
 segmenter = None
 
-# ================================
-# تحميل النموذج (من Google Drive)
-# ================================
+# =========================================================
+# تحميل النموذج
+# =========================================================
 def initialize_model():
+    """تحميل نموذج Unet EfficientNet من Google Drive"""
     global segmenter
+
     if not os.path.exists(MODEL_PATH):
         print("📥 تحميل النموذج من Google Drive...")
         gdown.download(MODEL_URL, MODEL_PATH, quiet=False)
 
     try:
-        print("🔄 تحميل DFUTissueSegNet...")
-        segmenter = DFUTissueSegNet(num_classes=3)
-        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+        print("🔄 تحميل Unet EfficientNet...")
+        model = smp.Unet(
+            encoder_name="efficientnet-b0",
+            encoder_weights=None,
+            classes=len(CLASS_NAMES),
+            activation="sigmoid"
+        )
 
-        # التعامل مع checkpoints التي تحتوي state_dict
-        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+        if "state_dict" in checkpoint:
             state_dict = checkpoint["state_dict"]
         else:
             state_dict = checkpoint
 
-        # تنظيف البوادئ الشائعة
-        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("✅ النموذج جاهز.")
+        clean_state = {k.replace("module.", "").replace("model.", ""): v for k, v in state_dict.items()}
+        model.load_state_dict(clean_state, strict=False)
+        model.to(DEVICE)
+        model.eval()
+        segmenter = model
+        print("✅ تم تحميل النموذج بنجاح.")
     except Exception as e:
         print(f"❌ فشل تحميل النموذج: {e}")
         import traceback; traceback.print_exc()
         segmenter = None
 
-# ================================
+# =========================================================
 # أدوات مساعدة
-# ================================
+# =========================================================
 def ensure_rgb(np_img):
+    """تحويل الصورة إلى RGB إذا لزم"""
     if np_img.ndim == 2:
         return cv2.cvtColor(np_img, cv2.COLOR_GRAY2RGB)
     if np_img.shape[-1] == 4:
         return cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
     return np_img
 
-def apply_legend_markdown():
-    return """
-### 🧭 مفتاح الألوان (Legend)
-- 🩸 **أحمر** → نسيج قرحة (Granulation)  
-- 🟡 **أصفر** → نسيج ميت جزئيًا (Slough)  
-- ⚫ **أسود** → نسيج نخر (Necrotic)  
-"""
-
-# ================================
-# التجزئة + الحساب + التلوين
-# ================================
-def segment_and_color(pil_img: Image.Image):
-    """
-    يُرجع:
-    - blended: الصورة مدموج عليها القناع اللوني
-    - mask_rgb: القناع اللوني (RGB)
-    - stats: نسب كل فئة + الإجمالي + مستوى الخطورة
-    """
-    if segmenter is None:
-        return pil_img, pil_img, {"خطأ": "النموذج غير مهيأ"}
-
-    # 1) التحضير
-    img_np = ensure_rgb(np.array(pil_img))
-    h, w = img_np.shape[:2]
-
-    # 2) التحجيم + التطبيع كما في التدريب
+def preprocess_image(img: Image.Image):
+    """تجهيز الصورة للنموذج"""
+    img_np = ensure_rgb(np.array(img))
     img_resized = cv2.resize(img_np, (IMG_SIZE, IMG_SIZE))
     img_norm = img_resized.astype(np.float32) / 255.0
-    img_norm = (img_norm - 0.5) / 0.5  # (x - 0.5) / 0.5
-
-    tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
-
-    # 3) التنبؤ
-    with torch.no_grad():
-        probs = segmenter(tensor).cpu().squeeze(0).numpy()  # (3, H, W) احتمالات
-
-    # 4) أقنعة ثنائية لكل فئة
-    masks = (probs >= THRESHOLD).astype(np.uint8)         # 0/1
-    # إزالة الضوضاء البسيطة
-    kernel = np.ones((5, 5), np.uint8)
-    for i in range(masks.shape[0]):
-        masks[i] = cv2.morphologyEx(masks[i], cv2.MORPH_OPEN, kernel)
-        masks[i] = cv2.morphologyEx(masks[i], cv2.MORPH_CLOSE, kernel)
-
-    # 5) حساب النسب على أبعاد الإدخال ثم إعادة القياس
-    total_pixels_input = IMG_SIZE * IMG_SIZE
-    ratios = {
-        CLASS_NAMES[0]: np.sum(masks[0]) / total_pixels_input * 100,
-        CLASS_NAMES[1]: np.sum(masks[1]) / total_pixels_input * 100,
-        CLASS_NAMES[2]: np.sum(masks[2]) / total_pixels_input * 100,
-    }
-    total_ratio = sum(ratios.values())
-
-    # 6) إنشاء قناع لوني على حجم الإدخال ثم إعادته لحجم الصورة الأصلي
-    color_mask = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)
-    color_mask[masks[0] == 1] = CLASS_COLORS[CLASS_NAMES[0]]
-    color_mask[masks[1] == 1] = CLASS_COLORS[CLASS_NAMES[1]]
-    color_mask[masks[2] == 1] = CLASS_COLORS[CLASS_NAMES[2]]
-
-    mask_rgb = cv2.resize(color_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    img_norm = (img_norm - 0.5) / 0.5
+    tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0)
+    return tensor.to(DEVICE), img_np
+
+# =========================================================
+# التجزئة والتحليل
+# =========================================================
+def analyze_image(img: Image.Image):
+    """تحليل صورة القدم وعرض النسب"""
+    if segmenter is None:
+        return img, img, {"خطأ": "النموذج غير مهيأ بعد."}
 
-    # 7) دمج القناع مع الصورة (ألفا ~0.5)
-    alpha = (cv2.cvtColor(mask_rgb, cv2.COLOR_RGB2GRAY) > 0).astype(np.float32)[..., None] * 0.5
-    blended = (alpha * mask_rgb + (1 - alpha) * img_np).astype(np.uint8)
+    try:
+        print("🔍 بدء التحليل...")
+        tensor, img_np = preprocess_image(img)
+
+        with torch.no_grad():
+            output = segmenter(tensor).cpu().squeeze(0).numpy()  # (3,H,W)
+
+        masks = (output >= THRESHOLD).astype(np.uint8)
+
+        # تنظيف الأقنعة
+        kernel = np.ones((5,5), np.uint8)
+        for i in range(masks.shape[0]):
+            masks[i] = cv2.morphologyEx(masks[i], cv2.MORPH_OPEN, kernel)
+            masks[i] = cv2.morphologyEx(masks[i], cv2.MORPH_CLOSE, kernel)
+
+        # حساب النسب
+        total_pixels = masks.shape[1] * masks.shape[2]
+        ratios = {
+            CLASS_NAMES[0]: np.sum(masks[0]) / total_pixels * 100,
+            CLASS_NAMES[1]: np.sum(masks[1]) / total_pixels * 100,
+            CLASS_NAMES[2]: np.sum(masks[2]) / total_pixels * 100
+        }
+        total_ratio = sum(ratios.values())
+
+        # إنشاء قناع لوني
+        color_mask = np.zeros((masks.shape[1], masks.shape[2], 3), dtype=np.uint8)
+        color_mask[masks[0] == 1] = CLASS_COLORS[CLASS_NAMES[0]]
+        color_mask[masks[1] == 1] = CLASS_COLORS[CLASS_NAMES[1]]
+        color_mask[masks[2] == 1] = CLASS_COLORS[CLASS_NAMES[2]]
+
+        color_mask = cv2.resize(color_mask, (img_np.shape[1], img_np.shape[0]))
+
+        # دمج القناع مع الصورة
+        alpha = 0.5
+        blended = cv2.addWeighted(img_np, 1 - alpha, color_mask, alpha, 0)
+
+        # تقييم الخطورة
+        if total_ratio == 0:
+            risk = "No Risk 🟢"
+        elif total_ratio < 1:
+            risk = "Low Risk 🟡"
+        elif total_ratio < 5:
+            risk = "Medium Risk 🟠"
+        else:
+            risk = "High Risk 🔴"
 
-    # 8) مستوى الخطورة
-    if total_ratio == 0:
-        risk = "No Risk 🟢"
-    elif total_ratio <= 1:
-        risk = "Low Risk 🟡"
-    elif total_ratio <= 5:
-        risk = "Medium Risk 🟠"
-    else:
-        risk = "High Risk 🔴"
+        report = {
+            "نسب الأنسجة (%)": {k: f"{v:.2f}" for k, v in ratios.items()},
+            "إجمالي (%)": f"{total_ratio:.2f}",
+            "مستوى الخطورة": risk
+        }
 
-    stats = {
-        "نِسَب_الأنسجة": {
-            CLASS_NAMES[0]: f"{ratios[CLASS_NAMES[0]]:.2f}%",
-            CLASS_NAMES[1]: f"{ratios[CLASS_NAMES[1]]:.2f}%",
-            CLASS_NAMES[2]: f"{ratios[CLASS_NAMES[2]]:.2f}%",
-            "الإجمالي": f"{total_ratio:.2f}%",
-        },
-        "مستوى_الخطورة": risk,
-        "ملاحظات": "التحليل يعتمد على DFUTissueSegNet متعدد الفئات (حجم 512 وتطبيع (x-0.5)/0.5)."
-    }
+        print(f"📊 النتائج: {report}")
+        return Image.fromarray(blended), Image.fromarray(color_mask), report
 
-    return Image.fromarray(blended), Image.fromarray(mask_rgb), stats
+    except Exception as e:
+        print(f"❌ خطأ أثناء التحليل: {e}")
+        import traceback; traceback.print_exc()
+        return img, img, {"خطأ": str(e)}
 
-# ================================
+# =========================================================
 # واجهة Gradio
-# ================================
+# =========================================================
 def build_ui():
-    with gr.Blocks(title="تحليل قرحة القدم - DFUTissueSegNet", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🦶 نظام تحليل قرحة القدم السكري - صورة واحدة")
-        gr.Markdown("يعتمد على **DFUTissueSegNet** لتجزئة الأنسجة وحساب نسبها، مع تلوين واضح وLegend.")
+    with gr.Blocks(title="تحليل قرحة القدم - EfficientNet Unet", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🦶 تحليل صورة القدم السكري (Unet + EfficientNet)")
+        gr.Markdown("الكشف عن أنواع الأنسجة المصابة (قرحة / Slough / نخر) وتقدير مستوى الخطورة.")
 
         with gr.Row():
             with gr.Column(scale=1):
                 input_img = gr.Image(type="pil", label="📤 ارفع صورة القدم", height=320)
                 analyze_btn = gr.Button("🔍 بدء التحليل", variant="primary")
-                legend = gr.Markdown(apply_legend_markdown())
 
             with gr.Column(scale=1):
                 out_blended = gr.Image(type="pil", label="🩸 الصورة مع القناع", height=320)
@@ -246,19 +175,17 @@ def build_ui():
                 out_json = gr.JSON(label="📊 التقرير التفصيلي")
 
         analyze_btn.click(
-            fn=segment_and_color,
+            fn=analyze_image,
             inputs=[input_img],
             outputs=[out_blended, out_mask, out_json]
         )
     return demo
 
-# ================================
+# =========================================================
 # تشغيل التطبيق
-# ================================
+# =========================================================
 if __name__ == "__main__":
     print("🚀 تهيئة النموذج...")
     initialize_model()
     app = build_ui()
-    # ملاحظة: على Spaces لا حاجة لـ share=True
-    app.launch(server_name="0.0.0.0", server_port=7860, share=False)
-
+    app.launch(server_name="0.0.0.0", server_port=7860)