Update app.py

app.py

@@ -1,12 +1,15 @@
 import gradio as gr
 import torch
 import random
-import requests
-import os
 from torchvision import transforms
 from PIL import Image
+import urllib.request
+import os
 
+# اسم الريبو الخاص بك على Hugging Face
 REPO_ID = "Alhdrawi/x_alhdrawi"
+
+# أسماء الملفات المرفوعة
 MODEL_FILES = [
     "best_128_0.0002_original_15000_0.859.pt",
     "best_128_0.0002_original_8000_0.857.pt",
@@ -20,44 +23,64 @@ MODEL_FILES = [
     "best_64_5e-05_original_22000_0.864.pt",
 ]
 
+# نفس المعالجة المستخدمة
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485], std=[0.229])
+])
+
+# الأمراض اللي يشخصها
 diseases = [
     "Atelectasis", "Cardiomegaly", "Consolidation", "Edema", "Effusion",
     "Emphysema", "Fibrosis", "Hernia", "Infiltration", "Mass", "Nodule",
     "Pleural_Thickening", "Pneumonia", "Pneumothorax"
 ]
 
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485], std=[0.229])
-])
+# تعرّف كلاس النموذج - لازم تعدله حسب المعمارية
+import torch.nn as nn
+
+class SimpleCNN(nn.Module):  # عدل هذا الكلاس حسب اللي دربت عليه
+    def __init__(self, num_classes=14):
+        super(SimpleCNN, self).__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 32, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1))
+        )
+        self.classifier = nn.Linear(64, num_classes)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = None
-current_model_path = None
+selected_model_file = random.choice(MODEL_FILES)
 
-def download_model_file(filename):
-    url = f"https://huggingface.co/{REPO_ID}/resolve/main/{filename}"
-    local_path = f"./{filename}"
+def download_and_load_model():
+    global model
+    url = f"https://huggingface.co/{REPO_ID}/resolve/main/{selected_model_file}"
+    local_path = f"/tmp/{selected_model_file}"
+    
     if not os.path.exists(local_path):
         print(f"Downloading model from {url}")
-        response = requests.get(url)
-        with open(local_path, "wb") as f:
-            f.write(response.content)
-    return local_path
-
-def load_model(_):
-    global model, current_model_path
-    selected_model = random.choice(MODEL_FILES)
-    local_path = download_model_file(selected_model)
-    model = torch.load(local_path, map_location=device)
+        urllib.request.urlretrieve(url, local_path)
+    
+    model = SimpleCNN(num_classes=len(diseases)).to(device)
+    state_dict = torch.load(local_path, map_location=device)
+    model.load_state_dict(state_dict)
     model.eval()
-    current_model_path = selected_model
-    return f"تم تحميل النموذج العشوائي: {selected_model}"
+    print(f"✅ Model loaded: {selected_model_file}")
 
 def predict(image):
     if model is None:
-        return "❗ الرجاء تحميل نموذج أولاً"
+        return "النموذج غير محمّل"
     img = transform(image.convert("L")).unsqueeze(0).to(device)
     with torch.no_grad():
         outputs = model(img)
@@ -65,16 +88,15 @@ def predict(image):
     results = {d: round(float(p), 3) for d, p in zip(diseases, probs)}
     return results
 
+# تحميل النموذج تلقائيًا عند بداية التشغيل
+download_and_load_model()
+
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 CheXzero | اختر صورة أشعة لتحليلها عبر نموذج يتم تحميله عشوائيًا")
-    with gr.Row():
-        load_button = gr.Button("تحميل نموذج عشوائي")
-        load_status = gr.Textbox(label="الحالة", interactive=False)
+    gr.Markdown(f"## 🧠 CheXzero | النموذج العشوائي المحمّل: `{selected_model_file}`")
     with gr.Row():
         image_input = gr.Image(type="pil", label="صورة الأشعة")
         output = gr.Label(num_top_classes=5)
 
-    load_button.click(fn=load_model, inputs=None, outputs=load_status)
     image_input.change(fn=predict, inputs=image_input, outputs=output)
 
 demo.launch()