Update app.py

app.py

@@ -1,82 +1,66 @@
 import torch
-import torch.nn as nn
-import torchvision.transforms as transforms
-from PIL import Image
-import numpy as np
+import torch.nn.functional as F
 import gradio as gr
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+
+# استيراد النموذج الحقيقي
+from model import DenseShuffleGCANet, extract_handcrafted_features
 
-# -------------------------------
-# 1️⃣ إعداد الفئات
-# -------------------------------
+# الفئات
 CLASSES = ["NONE", "INFECTION", "ISCHAEMIA", "BOTH"]
 
-# -------------------------------
-# 2️⃣ تعريف نموذج DenseShuffleGCANet
-# -------------------------------
-class DenseShuffleGCANet(nn.Module):
-    def __init__(self, num_classes=4, handcrafted_feature_dim=41):
-        super(DenseShuffleGCANet, self).__init__()
-        # مثال على backbone، عدلي حسب الكود الأصلي
-        self.backbone = nn.Sequential(
-            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-            nn.AdaptiveAvgPool2d((1,1))
-        )
-        self.fc_handcrafted = nn.Linear(handcrafted_feature_dim, 32)
-        self.classifier = nn.Linear(64 + 32, num_classes)
-    
-    def forward(self, x_image, x_features):
-        x_img = self.backbone(x_image)
-        x_img = x_img.view(x_img.size(0), -1)
-        x_feat = self.fc_handcrafted(x_features)
-        x = torch.cat([x_img, x_feat], dim=1)
-        out = self.classifier(x)
-        return out
+# الجهاز
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# -------------------------------
-# 3️⃣ تحميل النموذج مع weights
-# -------------------------------
+# تحميل النموذج
 model = DenseShuffleGCANet(num_classes=4, handcrafted_feature_dim=41)
-model.load_state_dict(torch.load("best_model_2.pth", map_location=torch.device('cpu')))
+model.load_state_dict(
+    torch.load("best_model.pth", map_location=device)
+)
+model.to(device)
 model.eval()
 
-# -------------------------------
-# 4️⃣ دالة استخراج الخصائص اليدوية
-# -------------------------------
-def extract_handcrafted_features(image_array):
-    features = np.random.rand(41).astype(np.float32)  # عدلي حسب خصائصك الحقيقية
-    return torch.tensor(features)
+# تحويل الصورة
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# دالة التنبؤ
+def predict(image: Image.Image):
+    image_tensor = transform(image).unsqueeze(0).to(device)
 
-# -------------------------------
-# 5️⃣ دالة التنبؤ
-# -------------------------------
-def predict_image(image: Image.Image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
-    ])
-    image_tensor = transform(image).unsqueeze(0)
-    features = extract_handcrafted_features(np.array(image)).unsqueeze(0)
+    features = extract_handcrafted_features(np.array(image))
+    features = features.unsqueeze(0).to(device)
 
     with torch.no_grad():
         outputs = model(image_tensor, features)
-        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
-        pred_class = CLASSES[probs.argmax()]
+        probs = F.softmax(outputs, dim=1)[0]
 
-    return {CLASSES[i]: float(probs[i]) for i in range(len(CLASSES))}, pred_class
+    result = {CLASSES[i]: float(probs[i]) for i in range(4)}
+    predicted_class = CLASSES[int(torch.argmax(probs))]
 
-# -------------------------------
-# 6️⃣ واجهة Gradio
-# -------------------------------
+    return result, predicted_class
+
+# واجهة Gradio
 interface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Label(num_top_classes=4), gr.Textbox(label="Predicted Class")],
-    title="DFU Foot Ulcer Classifier",
-    description="Upload an image of a foot ulcer to classify it as NONE, INFECTION, ISCHAEMIA, or BOTH."
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload DFU Image"),
+    outputs=[
+        gr.Label(num_top_classes=4, label="Probabilities"),
+        gr.Textbox(label="Predicted Class")
+    ],
+    title="DFU Classification System",
+    description="Classifies diabetic foot images into NONE, INFECTION, ISCHAEMIA, or BOTH."
 )
 
 if __name__ == "__main__":
     interface.launch()
 
+