Update app.py

app.py

@@ -1,24 +1,46 @@
 # -*- coding: utf-8 -*-
+# -*- coding: utf-8 -*-
+"""
+واجهة موحّدة لتحليل صورة القدم باستخدام DFUTissueSegNet (من Google Drive)
+- ألوان متعددة: قرحة=أحمر، Slough=أصفر، نخر=أسود
+- حساب نسب كل نوع + مستوى خطورة إجمالي
+- Legend داخل Gradio
+"""
+
+import os
+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 cv2
-import numpy as np
+
 import gradio as gr
-import gdown
-import os
-from PIL import Image
 
-# =========================================================
-# الإعدادات العامة
-# =========================================================
+# ================================
+# إعدادات عامة
+# ================================
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-IMG_SIZE = 512  # حجم الإدخال الصحيح للنموذج
-CLASS_NAMES = ["Ulcer", "Slough", "Necrosis"]
-
-# =========================================================
-# تعريف نموذج DFUTissueSegNet
-# =========================================================
+IMG_SIZE = 512                    # حجم الإدخال كما في تدريب DFUTissueSegNet
+THRESHOLD = 0.4                   # عتبة تحويل الاحتمال إلى قناع (جرّبي 0.35..0.6 حسب بياناتك)
+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_COLORS = {
+    "قرحة (Granulation)": (255, 0, 0),     # أحمر
+    "Slough (أنسجة ميتة جزئيًا)": (255, 255, 0),  # أصفر
+    "نخر (Necrotic)": (0, 0, 0)            # أسود
+}
+
+# ================================
+# نموذج DFUTissueSegNet
+# ================================
 class ConvBlock(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
@@ -28,23 +50,18 @@ class ConvBlock(nn.Module):
             nn.ReLU(inplace=True),
             nn.Conv2d(out_ch, out_ch, 3, padding=1),
             nn.BatchNorm2d(out_ch),
-            nn.ReLU(inplace=True)
+            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.encoder1 = ConvBlock(3, 64)
-        self.pool1 = nn.MaxPool2d(2)
-        self.encoder2 = ConvBlock(64, 128)
-        self.pool2 = nn.MaxPool2d(2)
-        self.encoder3 = ConvBlock(128, 256)
-        self.pool3 = nn.MaxPool2d(2)
-        self.encoder4 = ConvBlock(256, 512)
-        self.pool4 = nn.MaxPool2d(2)
+        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)
 
@@ -60,43 +77,29 @@ class DFUTissueSegNet(nn.Module):
         self.final = nn.Conv2d(64, num_classes, 1)
 
     def forward(self, x):
-        e1 = self.encoder1(x)
-        e2 = self.encoder2(self.pool1(e1))
-        e3 = self.encoder3(self.pool2(e2))
-        e4 = self.encoder4(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)
-
+        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))
 
-# =========================================================
-# تحميل النموذج من Google Drive
-# =========================================================
 segmenter = None
 
+# ================================
+# تحميل النموذج (من Google Drive)
+# ================================
 def initialize_model():
-    """تحميل DFUTissueSegNet من Google Drive مع دعم checkpoint"""
     global segmenter
-    MODEL_URL = "https://drive.google.com/uc?id=1Ovaczsjdp3E-_gYF2pbUibDjPWAC1a6c"
-    MODEL_PATH = "best_model_5.pth"
-
     if not os.path.exists(MODEL_PATH):
-        print("📥 تحميل DFUTissueSegNet من Google Drive...")
+        print("📥 تحميل النموذج من Google Drive...")
         gdown.download(MODEL_URL, MODEL_PATH, quiet=False)
 
     try:
@@ -104,12 +107,13 @@ def initialize_model():
         segmenter = DFUTissueSegNet(num_classes=3)
         checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
 
-        # إذا كان checkpoint يحتوي على state_dict
-        if "state_dict" in checkpoint:
+        # التعامل مع checkpoints التي تحتوي state_dict
+        if isinstance(checkpoint, dict) and "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.", "")
@@ -118,111 +122,143 @@ def initialize_model():
         segmenter.load_state_dict(clean_state, strict=False)
         segmenter.to(DEVICE)
         segmenter.eval()
-        print("✅ تم تحميل DFUTissueSegNet من Google Drive بنجاح.")
+        print("✅ النموذج جاهز.")
     except Exception as e:
         print(f"❌ فشل تحميل النموذج: {e}")
-        import traceback
-        traceback.print_exc()
+        import traceback; traceback.print_exc()
         segmenter = None
 
-# =========================================================
-# دوال التحليل
-# =========================================================
-def analyze_image(img: Image.Image):
-    """تحليل صورة القدم وإخراج القناع"""
+# ================================
+# أدوات مساعدة
+# ================================
+def ensure_rgb(np_img):
+    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 img, {"خطأ": "النموذج غير مهيأ بعد."}
-
-    try:
-        print("🔍 بدء التحليل...")
-        img_np = np.array(img)
-
-        # التأكد من أن الصورة RGB
-        if len(img_np.shape) == 2:
-            img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
-        elif img_np.shape[-1] == 4:
-            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
-
-        # التحجيم والتطبيع الصحيح
-        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  # كما في التدريب
-
-        tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
-
-        with torch.no_grad():
-            output = segmenter(tensor).cpu().squeeze(0).numpy()
-
-        # تطبيق العتبة لكل قناة
-        th = 0.4
-        masks = (output >= th).astype(np.uint8)
-
-        # حساب النسب
-        ulcer_ratio = np.sum(masks[0]) / (IMG_SIZE * IMG_SIZE) * 100
-        slough_ratio = np.sum(masks[1]) / (IMG_SIZE * IMG_SIZE) * 100
-        necrosis_ratio = np.sum(masks[2]) / (IMG_SIZE * IMG_SIZE) * 100
-        total_ratio = ulcer_ratio + slough_ratio + necrosis_ratio
-
-        # تلوين القناع
-        color_mask = np.zeros_like(img_resized)
-        color_mask[masks[0] == 1] = [255, 0, 0]      # قرحة حمراء
-        color_mask[masks[1] == 1] = [255, 255, 0]    # Slough صفراء
-        color_mask[masks[2] == 1] = [0, 0, 0]        # نخر أسود
-
-        # دمج القناع مع الصورة
-        blended = cv2.addWeighted(img_resized, 0.7, color_mask, 0.6, 0)
-        blended = cv2.cvtColor(blended, cv2.COLOR_RGB2BGR)
-
-        # تقييم الخطورة
-        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 = {
-            "قرحة (%)": round(ulcer_ratio, 2),
-            "Slough (%)": round(slough_ratio, 2),
-            "نخر (%)": round(necrosis_ratio, 2),
-            "إجمالي (%)": round(total_ratio, 2),
-            "مستوى الخطورة": risk
-        }
-
-        print(f"📊 تقرير التحليل: {report}")
-        return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)), report
-
-    except Exception as e:
-        print(f"❌ خطأ أثناء التحليل: {e}")
-        import traceback
-        traceback.print_exc()
-        return img, {"خطأ": str(e)}
-
-# =========================================================
+        return pil_img, pil_img, {"خطأ": "النموذج غير مهيأ"}
+
+    # 1) التحضير
+    img_np = ensure_rgb(np.array(pil_img))
+    h, w = img_np.shape[:2]
+
+    # 2) التحجيم + التطبيع كما في التدريب
+    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)
+
+    # 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)
+
+    # 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 🔴"
+
+    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)."
+    }
+
+    return Image.fromarray(blended), Image.fromarray(mask_rgb), stats
+
+# ================================
 # واجهة Gradio
-# =========================================================
-print("🚀 تهيئة نموذج DFUTissueSegNet...")
-initialize_model()
-
-with gr.Blocks(title="تحليل DFUTissueSegNet") as demo:
-    gr.Markdown("""
-    # 🩸 نظام تحليل DFU Tissue Segmentation
-    **نموذج DFUTissueSegNet** لتجزئة أنسجة القدم السكري وتقدير مستوى الخطورة.
-    """)
-
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(type="pil", label="📤 ارفع صورة القدم")
-            analyze_btn = gr.Button("🔍 تحليل الصورة", variant="primary")
-
-        with gr.Column():
-            output_img = gr.Image(type="pil", label="🩸 النتيجة", height=400)
-            output_json = gr.JSON(label="📊 تقرير التحليل")
-
-    analyze_btn.click(fn=analyze_image, inputs=[input_img], outputs=[output_img, output_json])
+# ================================
+def build_ui():
+    with gr.Blocks(title="تحليل قرحة القدم - DFUTissueSegNet", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🦶 نظام تحليل قرحة القدم السكري - صورة واحدة")
+        gr.Markdown("يعتمد على **DFUTissueSegNet** لتجزئة الأنسجة وحساب نسبها، مع تلوين واضح وLegend.")
+
+        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)
+                out_mask = gr.Image(type="pil", label="🧩 القناع اللوني", height=320)
+                out_json = gr.JSON(label="📊 التقرير التفصيلي")
+
+        analyze_btn.click(
+            fn=segment_and_color,
+            inputs=[input_img],
+            outputs=[out_blended, out_mask, out_json]
+        )
+    return demo
 
+# ================================
+# تشغيل التطبيق
+# ================================
 if __name__ == "__main__":
-    print("🌐 بدء التشغيل...")
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    print("🚀 تهيئة النموذج...")
+    initialize_model()
+    app = build_ui()
+    # ملاحظة: على Spaces لا حاجة لـ share=True
+    app.launch(server_name="0.0.0.0", server_port=7860, share=False)
+