Update app.py

app.py

@@ -1,57 +1,50 @@
 # -*- coding: utf-8 -*-
-"""
-نظام تحليل قرحة القدم السكري باستخدام DFUTissueSegNet
-- يعتمد فقط على التجزئة متعددة الفئات (قرحة / Slough / نخر)
-- يعرض الألوان + نسب كل نوع نسيج + مستوى الخطورة
-"""
-
-import os
-import numpy as np
-from PIL import Image
-import cv2
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import gdown
+import cv2
+import numpy as np
 import gradio as gr
+import gdown
+import os
+from PIL import Image
 
-
-# ======================================================
+# =========================================================
 # الإعدادات العامة
-# ======================================================
-segmenter = None
-IMG_SIZE = 224
+# =========================================================
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+IMG_SIZE = 512  # حجم الإدخال الصحيح للنموذج
+CLASS_NAMES = ["Ulcer", "Slough", "Necrosis"]
 
-
-# ======================================================
-# تعريف النموذج DFUTissueSegNet
-# ======================================================
+# =========================================================
+# تعريف نموذج DFUTissueSegNet
+# =========================================================
 class ConvBlock(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(ConvBlock, self).__init__()
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
         self.block = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 3, padding=1),
-            nn.BatchNorm2d(out_channels),
+            nn.Conv2d(in_ch, out_ch, 3, padding=1),
+            nn.BatchNorm2d(out_ch),
             nn.ReLU(inplace=True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1),
-            nn.BatchNorm2d(out_channels),
+            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(DFUTissueSegNet, self).__init__()
+        super().__init__()
         self.encoder1 = ConvBlock(3, 64)
-        self.pool1 = nn.MaxPool2d(2, 2)
+        self.pool1 = nn.MaxPool2d(2)
         self.encoder2 = ConvBlock(64, 128)
-        self.pool2 = nn.MaxPool2d(2, 2)
+        self.pool2 = nn.MaxPool2d(2)
         self.encoder3 = ConvBlock(128, 256)
-        self.pool3 = nn.MaxPool2d(2, 2)
+        self.pool3 = nn.MaxPool2d(2)
         self.encoder4 = ConvBlock(256, 512)
-        self.pool4 = nn.MaxPool2d(2, 2)
+        self.pool4 = nn.MaxPool2d(2)
 
         self.center = ConvBlock(512, 1024)
 
@@ -71,9 +64,9 @@ class DFUTissueSegNet(nn.Module):
         e2 = self.encoder2(self.pool1(e1))
         e3 = self.encoder3(self.pool2(e2))
         e4 = self.encoder4(self.pool3(e3))
-        center = self.center(self.pool4(e4))
+        c = self.center(self.pool4(e4))
 
-        d4 = self.up4(center)
+        d4 = self.up4(c)
         d4 = torch.cat([d4, e4], dim=1)
         d4 = self.dec4(d4)
 
@@ -89,16 +82,15 @@ class DFUTissueSegNet(nn.Module):
         d1 = torch.cat([d1, e1], dim=1)
         d1 = self.dec1(d1)
 
-        out = self.final(d1)
-        out = F.softmax(out, dim=1)
-        return out
+        return torch.sigmoid(self.final(d1))
 
+# =========================================================
+# تحميل النموذج من Google Drive
+# =========================================================
+segmenter = None
 
-# ======================================================
-# تحميل النموذج
-# ======================================================
 def initialize_model():
-    """تحميل DFUTissueSegNet من Google Drive مع دعم checkpoint الكامل"""
+    """تحميل DFUTissueSegNet من Google Drive مع دعم checkpoint"""
     global segmenter
     MODEL_URL = "https://drive.google.com/uc?id=1Ovaczsjdp3E-_gYF2pbUibDjPWAC1a6c"
     MODEL_PATH = "best_model_5.pth"
@@ -112,13 +104,12 @@ def initialize_model():
         segmenter = DFUTissueSegNet(num_classes=3)
         checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
 
-        # بعض النماذج تكون محفوظة داخل "state_dict"
+        # إذا كان checkpoint يحتوي على 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.", "")
@@ -134,182 +125,104 @@ def initialize_model():
         traceback.print_exc()
         segmenter = None
 
-
-
-# ======================================================
-# دالة التجزئة
-# ======================================================
-def segment_ulcer(pil_img: Image.Image):
-    """تجزئة متعددة الفئات + حساب نسب كل نوع نسيج"""
+# =========================================================
+# دوال التحليل
+# =========================================================
+def analyze_image(img: Image.Image):
+    """تحليل صورة القدم وإخراج القناع"""
     if segmenter is None:
-        np_img = np.array(pil_img)
-        return np.zeros((np_img.shape[0], np_img.shape[1], 3), dtype=np.uint8), {}
+        return img, {"خطأ": "النموذج غير مهيأ بعد."}
 
     try:
-        img_np = np.array(pil_img)
-        if img_np.ndim == 2:
+        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)
-            pred = output.squeeze().cpu().numpy()  # (3, H, W)
-
-        # تأكيد الأبعاد
-        if pred.ndim != 3 or pred.shape[0] < 3:
-            print("⚠️ النموذج لم يُرجع 3 قنوات. سيتم استخدام القناة الأولى فقط.")
-            pred = np.stack([pred, np.zeros_like(pred), np.zeros_like(pred)], axis=0)
-
-        pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
-
-        gran = pred[0, :, :]  # أحمر
-        slough = pred[1, :, :]  # أصفر
-        nec = pred[2, :, :]  # أسود
-
-        th = 0.55
-        gran_mask = (gran >= th).astype(np.uint8)
-        slough_mask = (slough >= th).astype(np.uint8)
-        nec_mask = (nec >= th).astype(np.uint8)
-
-        kernel = np.ones((5, 5), np.uint8)
-        for m in [gran_mask, slough_mask, nec_mask]:
-            cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel)
-            cv2.morphologyEx(m, cv2.MORPH_CLOSE, kernel)
-
-        mask_color = np.zeros((*gran_mask.shape, 3), dtype=np.uint8)
-        mask_color[gran_mask == 1] = (255, 0, 0)
-        mask_color[slough_mask == 1] = (255, 255, 0)
-        mask_color[nec_mask == 1] = (0, 0, 0)
-
-        mask_resized = cv2.resize(mask_color, (img_np.shape[1], img_np.shape[0]), interpolation=cv2.INTER_NEAREST)
-
-        total_pixels = mask_resized.shape[0] * mask_resized.shape[1]
-        gran_ratio = np.sum(gran_mask) / total_pixels * 100
-        slough_ratio = np.sum(slough_mask) / total_pixels * 100
-        nec_ratio = np.sum(nec_mask) / total_pixels * 100
-        total_ratio = gran_ratio + slough_ratio + nec_ratio
-
-        tissue_stats = {
-            "قرحة (Granulation)": f"{gran_ratio:.2f}%",
-            "Slough (أنسجة ميتة جزئيًا)": f"{slough_ratio:.2f}%",
-            "نخر (Necrotic)": f"{nec_ratio:.2f}%",
-            "الإجمالي": f"{total_ratio:.2f}%"
-        }
+            output = segmenter(tensor).cpu().squeeze(0).numpy()
 
-        return mask_resized, tissue_stats
-
-    except Exception as e:
-        print(f"❌ خطأ في التجزئة: {e}")
-        import traceback
-        traceback.print_exc()
-        np_img = np.array(pil_img)
-        return np.zeros((np_img.shape[0], np_img.shape[1], 3), dtype=np.uint8), {}
+        # تطبيق العتبة لكل قناة
+        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
 
-# ======================================================
-# تطبيق القناع اللوني
-# ======================================================
-def apply_ulcer_mask(pil_img: Image.Image, mask_color):
-    try:
-        base = np.array(pil_img)
-        if base.ndim == 2:
-            base = cv2.cvtColor(base, cv2.COLOR_GRAY2RGB)
-        elif base.shape[-1] == 4:
-            base = cv2.cvtColor(base, cv2.COLOR_RGBA2RGB)
-
-        if mask_color.shape[:2] != base.shape[:2]:
-            mask_color = cv2.resize(mask_color, (base.shape[1], base.shape[0]))
-
-        mask_gray = cv2.cvtColor(mask_color, cv2.COLOR_RGB2GRAY)
-        alpha = (mask_gray > 0).astype(np.float32)[..., None] * 0.5
-        blended = (alpha * mask_color + (1 - alpha) * base).astype(np.uint8)
-        return Image.fromarray(blended)
-    except Exception as e:
-        print(f"❌ خطأ في تطبيق القناع: {e}")
-        return pil_img
+        # تلوين القناع
+        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)
 
-# ======================================================
-# تحليل الصورة
-# ======================================================
-def analyze_single_image(pil_img: Image.Image):
-    if pil_img is None:
-        return None, None, {"خطأ": "لم يتم رفع صورة"}
-
-    try:
-        mask_color, tissue_stats = segment_ulcer(pil_img)
-        blended = apply_ulcer_mask(pil_img, mask_color)
-
-        total_ratio = float(tissue_stats.get("الإجمالي", "0").replace("%", ""))
+        # تقييم الخطورة
         if total_ratio == 0:
-            level, emoji = "No Risk", "🟢"
-        elif total_ratio <= 1:
-            level, emoji = "Low Risk", "🟡"
-        elif total_ratio <= 5:
-            level, emoji = "Medium Risk", "🟠"
+            risk = "No Risk 🟢"
+        elif total_ratio < 1:
+            risk = "Low Risk 🟡"
+        elif total_ratio < 5:
+            risk = "Medium Risk 🟠"
         else:
-            level, emoji = "High Risk", "🔴"
+            risk = "High Risk 🔴"
 
         report = {
-            "مستوى_الخطورة": f"{level} {emoji}",
-            "نسب_الأنسجة": tissue_stats,
-            "ملاحظات": "تم التحليل اعتمادًا على DFUTissueSegNet متعدد الفئات."
+            "قرحة (%)": round(ulcer_ratio, 2),
+            "Slough (%)": round(slough_ratio, 2),
+            "نخر (%)": round(necrosis_ratio, 2),
+            "إجمالي (%)": round(total_ratio, 2),
+            "مستوى الخطورة": risk
         }
 
-        return blended, Image.fromarray(mask_color), report
+        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 pil_img, None, {"خطأ": str(e)}
-
+        return img, {"خطأ": str(e)}
 
-# ======================================================
+# =========================================================
 # واجهة Gradio
-# ======================================================
-def build_ui():
-    with gr.Blocks(title="تحليل قرحة القدم السكري", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🦶 نظام تحليل قرحة القدم السكري - DFUTissueSegNet")
-        gr.Markdown("### يعتمد التحليل على التجزئة لتحديد أنواع الأنسجة المتضررة ونسبة كل نوع")
-
-        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=320)
-                mask_img = gr.Image(type="pil", label="🧩 القناع اللوني", height=320)
-                json_out = gr.JSON(label="📊 التقرير التفصيلي")
-
-        gr.Markdown("""
-        ---
-        ### 🧭 مفتاح الألوان (Legend)
-        - 🩸 **أحمر** → نسيج قرحة (Granulation)  
-        - 🟡 **أصفر** → نسيج ميت جزئيًا (Slough)  
-        - ⚫ **أسود** → نسيج نخر (Necrotic)  
-        ---
-        """)
-
-        analyze_btn.click(
-            fn=analyze_single_image,
-            inputs=[input_img],
-            outputs=[output_img, mask_img, json_out]
-        )
-    return demo
+# =========================================================
+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])
 
-# ======================================================
-# ��شغيل النظام
-# ======================================================
 if __name__ == "__main__":
-    print("🚀 تهيئة نموذج DFUTissueSegNet...")
-    initialize_model()
-    demo = build_ui()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    print("🌐 بدء التشغيل...")
+    demo.launch(server_name="0.0.0.0", server_port=7860)