Update app.py

app.py

@@ -1,37 +1,39 @@
 """
 ╔══════════════════════════════════════════════════════════════╗
-║  🔱 SONAR-AI v17.2 — Swin-V2 Upgrade Patch                 ║
-║  أضف هذا الكود على v16.8 الأصلي                            ║
+║  🔱 swin_v2_addon.py — Swin-V2 Module for SONAR-AI         ║
+║  ✅ ضع هذا الملف بجانب app.py على HuggingFace              ║
 ╠══════════════════════════════════════════════════════════════╣
-║  ✅ يحافظ على كل شيء في v16.8 (Physics + PRMI + VTA)       ║
-║  ✅ يضيف Swin-V2 Classification (94.1% F1, 43 classes)      ║
-║  ✅ يضيف Swin-V2 Concealment Detection (98.9% F1)           ║
-║  ✅ يضيف Swin-V2 Risk Assessment (97.2% F1)                 ║
-║  ✅ يعرض نتائج Swin-V2 في صفحة التحليل                     ║
+║  في app.py الأصلي (v16.8) أضف فقط 3 أسطر:                 ║
+║                                                              ║
+║  1. في أعلى الملف:                                           ║
+║     from swin_v2_addon import *                              ║
+║                                                              ║
+║  2. في نهاية load_all_models() قبل MODELS_LOADED:           ║
+║     load_swinv2_models(CLASSIFICATION_MODELS, DEVICE)        ║
+║                                                              ║
+║  3. في analyze_image() قبل en_desc:                          ║
+║     swin_html = run_swinv2_analysis(img, DEVICE)             ║
+║     ثم أضف {swin_html} في HTML بعد {anomaly_html}           ║
 ╚══════════════════════════════════════════════════════════════╝
-
-التعليمات:
-────────────
-1. افتح app.py الأصلي (v16.8) في HuggingFace
-2. ابحث عن كل "=== ADD HERE ===" وأضف الكود في المكان المحدد
-3. غيّر VERSION من "16.8" إلى "17.2"
 """
 
-# ═══════════════════════════════════════════════════════════════
-# التعديل 1: غيّر الإصدار
-# ابحث عن: VERSION = "16.8"
-# استبدل بـ:
-# ═══════════════════════════════════════════════════════════════
+import numpy as np
 
-VERSION = "17.2"  # ← غيّر هذا السطر
+# Try importing torch
+try:
+    import torch
+    import torch.nn as nn
+    from torchvision import transforms
+    TORCH_OK = True
+    ADDON_DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+except ImportError:
+    TORCH_OK = False
+    ADDON_DEVICE = None
 
 # ═══════════════════════════════════════════════════════════════
-# التعديل 2: أضف SWINV2_CATEGORIES
-# ابحث عن: HS_CHAPTERS = {
-# أضف الكود التالي قبلها مباشرة:
+# 🏷️ Categories
 # ═══════════════════════════════════════════════════════════════
 
-# 🆕 فئات Swin-V2 المدربة (43 صنف — مرتبة أبجدياً)
 SWINV2_CATEGORIES = sorted([
     'appliances', 'auto_parts', 'bags', 'banana', 'batteries', 'beverages',
     'cables', 'canned_food', 'ceramic', 'chemicals', 'cleaning', 'clothes',
@@ -45,206 +47,237 @@ SWINV2_CATEGORIES = sorted([
 SWINV2_CONCEALMENT_CLASSES = ['match', 'no_match']
 SWINV2_RISK_CLASSES = ['critical', 'high', 'low', 'medium', 'safe']
 
-
 # ═══════════════════════════════════════════════════════════════
-# التعديل 3: أضف دالة transforms خاصة بـ Swin-V2
-# ابحث عن: def get_transforms():
-# أضف بعدها مباشرة:
+# Global model references
 # ═══════════════════════════════════════════════════════════════
 
-def get_swin_transforms():
-    """Transforms for Swin-V2 trained models (256×256)"""
-    if torch is None:
+SWINV2_CLASSIFICATION = None
+SWINV2_CONCEALMENT = None
+SWINV2_RISK = None
+
+
+def get_swinv2_transforms(img_size=256):
+    if not TORCH_OK:
         return None
     return transforms.Compose([
-        transforms.Resize((256, 256)),
+        transforms.Resize((img_size, img_size)),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
 
 
-# ═══════════════════════════════════════════════════════════════
-# التعديل 4: أضف تحميل Swin-V2 في load_all_models
-# ابحث عن: MODELS_LOADED = len(CLASSIFICATION_MODELS) > 0
-# أضف الكود التالي قبلها:
-# ═══════════════════════════════════════════════════════════════
+def find_swin_model(name):
+    """Find model file locally or download from HuggingFace"""
+    import os, glob
+    MODEL_REPO = "DrAbbas/SONAR-AI-Models"
+    
+    for base in [".", "/app", os.path.dirname(os.path.abspath(__file__))]:
+        path = os.path.join(base, name)
+        if os.path.exists(path):
+            return path
+    try:
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(repo_id=MODEL_REPO, filename=name)
+        return path
+    except:
+        pass
+    return None
 
-    # === ADD BEFORE "MODELS_LOADED = len(CLASSIFICATION_MODELS) > 0" ===
 
-    # 🆕 Swin-V2 Classification (43 classes, 94.1% F1)
-    SWIN_CLS_MODEL = None
-    path = find_model("swin_v2_cls_best.pt")
+def load_swinv2_models(classification_models_dict, device=None):
+    """
+    تحميل 3 نماذج Swin-V2.
+    استدعيها في load_all_models() في app.py:
+        load_swinv2_models(CLASSIFICATION_MODELS, DEVICE)
+    """
+    global SWINV2_CLASSIFICATION, SWINV2_CONCEALMENT, SWINV2_RISK
+    
+    if not TORCH_OK:
+        print("⚠️ Swin-V2: PyTorch not available")
+        return
+    
+    if device is None:
+        device = ADDON_DEVICE
+    
+    # 1. Swin-V2 Classification (43 classes, 94.1% F1)
+    path = find_swin_model("classification/best_swinv2_43cls.pth")
     if path:
         try:
             import timm
-            nc_swin = len(SWINV2_CATEGORIES)
-            m = timm.create_model('swinv2_tiny_window8_256', pretrained=False, num_classes=nc_swin)
-            checkpoint = torch.load(path, map_location=DEVICE)
-            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-                m.load_state_dict(checkpoint['model_state_dict'])
+            nc = len(SWINV2_CATEGORIES)
+            m = timm.create_model('swinv2_tiny_window8_256', pretrained=False, num_classes=nc)
+            ckpt = torch.load(path, map_location=device)
+            if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
+                m.load_state_dict(ckpt['model_state_dict'])
             else:
-                m.load_state_dict(checkpoint)
-            m.to(DEVICE).eval()
-            CLASSIFICATION_MODELS['Swin-V2'] = m
-            SWIN_CLS_MODEL = m
-            print(f"✅ Swin-V2 Classification loaded ({nc_swin}cls, 94.1% F1)")
+                m.load_state_dict(ckpt)
+            m.to(device).eval()
+            classification_models_dict['Swin-V2'] = m
+            SWINV2_CLASSIFICATION = m
+            print(f"✅ Swin-V2 Classification loaded ({nc}cls, 94.1% F1)")
         except Exception as e:
             print(f"⚠️ Swin-V2 Classification: {e}")
-
-    # 🆕 Swin-V2 Concealment Detection (98.9% F1)
-    global CONCEALMENT_MODEL
-    CONCEALMENT_MODEL = None
-    path = find_model("swin_v2_concealment_best.pt")
-    if path:
+    
+    # 2. Concealment Detection (98.9% F1)
+    conc_path = find_swin_model("concealment/best_eva02.pth") or find_swin_model("concealment/best_swinv2.pth")
+    if conc_path:
         try:
             import timm
-            m = timm.create_model('swinv2_tiny_window8_256', pretrained=False, num_classes=2)
-            checkpoint = torch.load(path, map_location=DEVICE)
-            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-                m.load_state_dict(checkpoint['model_state_dict'])
+            if 'eva02' in conc_path:
+                m = timm.create_model('eva02_tiny_patch14_224', pretrained=False, num_classes=2)
+                img_size_conc = 224
+            else:
+                m = timm.create_model('swinv2_tiny_window8_256', pretrained=False, num_classes=2)
+                img_size_conc = 256
+            ckpt = torch.load(conc_path, map_location=device)
+            if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
+                m.load_state_dict(ckpt['model_state_dict'])
             else:
-                m.load_state_dict(checkpoint)
-            m.to(DEVICE).eval()
-            CONCEALMENT_MODEL = m
-            print(f"✅ Swin-V2 Concealment loaded (98.9% F1)")
+                m.load_state_dict(ckpt)
+            m.to(device).eval()
+            m._conc_img_size = img_size_conc
+            SWINV2_CONCEALMENT = m
+            print(f"✅ Concealment loaded (98.9% F1)")
         except Exception as e:
-            print(f"⚠️ Swin-V2 Concealment: {e}")
-
-    # 🆕 Swin-V2 Risk Assessment (5 levels, 97.2% F1)
-    global RISK_ASSESSMENT_MODEL
-    RISK_ASSESSMENT_MODEL = None
-    path = find_model("swin_v2_risk_best.pt")
-    if path:
+            SWINV2_CONCEALMENT = None
+            print(f"⚠️ Concealment: {e}")
+    
+    # 3. Risk Assessment (5 levels, 97.2% F1)
+    risk_path = find_swin_model("risk/best_swinv2_risk.pth")
+    if risk_path:
         try:
             import timm
             m = timm.create_model('swinv2_tiny_window8_256', pretrained=False, num_classes=5)
-            checkpoint = torch.load(path, map_location=DEVICE)
-            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-                m.load_state_dict(checkpoint['model_state_dict'])
+            ckpt = torch.load(risk_path, map_location=device)
+            if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
+                m.load_state_dict(ckpt['model_state_dict'])
             else:
-                m.load_state_dict(checkpoint)
-            m.to(DEVICE).eval()
-            RISK_ASSESSMENT_MODEL = m
+                m.load_state_dict(ckpt)
+            m.to(device).eval()
+            SWINV2_RISK = m
             print(f"✅ Swin-V2 Risk loaded (5 levels, 97.2% F1)")
         except Exception as e:
+            SWINV2_RISK = None
             print(f"⚠️ Swin-V2 Risk: {e}")
+    
+    swin_count = sum(1 for x in [SWINV2_CLASSIFICATION, SWINV2_CONCEALMENT, SWINV2_RISK] if x is not None)
+    print(f"🔱 Swin-V2: {swin_count}/3 models loaded")
 
 
-# ═══════════════════════════════════════════════════════════════
-# التعديل 5: إصلاح classify_image لـ Swin-V2
-# ابحث عن "if name in ('YOLOv10x-cls', 'YOLO11x-cls'):" في دالة classify_image
-# استبدل الـ else (الذي يشتغل على نماذج PyTorch) بـ:
-# ═══════════════════════════════════════════════════════════════
-
-            # === REPLACE the "else:" block in classify_image with: ===
-            else:
-                # Choose correct transforms and categories
-                if name == 'Swin-V2':
-                    img_transform = get_swin_transforms()
-                    cats = SWINV2_CATEGORIES
-                else:
-                    img_transform = get_transforms()
-                    cats = CATEGORIES
-
-                input_t = img_transform(img).unsqueeze(0).to(DEVICE)
-                with torch.no_grad():
-                    out = model(input_t)
-                    probs = torch.softmax(out, dim=1)[0].cpu().numpy()
-                    for idx in range(len(probs)):
-                        if idx < len(cats):
-                            cat = cats[idx]
-                            all_preds.setdefault(cat, []).append(float(probs[idx]))
-                    top_idx = probs.argmax()
-                    if top_idx < len(cats):
-                        model_results[name] = {'top1': cats[top_idx], 'confidence': float(probs[top_idx])}
-
-
-# ═══════════════════════════════════════════════════════════════
-# التعديل 6: عرض Swin-V2 في analyze_image
-# ابحث عن "ar_desc = " في analyze_image
-# أضف بعدها الكود التالي لعرض Swin-V2 results:
-# ═══════════════════════════════════════════════════════════════
+def classify_swinv2(img, device=None):
+    """
+    تصنيف صورة بـ Swin-V2 — يُرجع (category, confidence)
+    """
+    if SWINV2_CLASSIFICATION is None or not TORCH_OK:
+        return None, 0.0
+    
+    if device is None:
+        device = ADDON_DEVICE
+    
+    try:
+        t = get_swinv2_transforms(256)
+        input_t = t(img.convert('RGB')).unsqueeze(0).to(device)
+        with torch.no_grad():
+            out = SWINV2_CLASSIFICATION(input_t)
+            probs = torch.softmax(out, dim=1)[0].cpu().numpy()
+            top_idx = probs.argmax()
+            if top_idx < len(SWINV2_CATEGORIES):
+                return SWINV2_CATEGORIES[top_idx], float(probs[top_idx])
+    except:
+        pass
+    return None, 0.0
 
-    # === ADD AFTER "ar_desc = ..." line ===
 
-    # 🆕 Swin-V2 Concealment + Risk Assessment
-    swin_concealment_html = ""
-    swin_risk_html = ""
+def run_swinv2_analysis(img, device=None):
+    """
+    تشغيل 3 نماذج Swin-V2 وإرجاع HTML.
+    استدعيها في analyze_image() في app.py:
+        swin_html = run_swinv2_analysis(img, DEVICE)
+    ثم أضف {swin_html} في HTML بعد {anomaly_html}
+    """
+    if img is None or not TORCH_OK:
+        return ""
+    
+    if device is None:
+        device = ADDON_DEVICE
+    
+    html = ""
+    
+    # Swin-V2 Classification
+    if SWINV2_CLASSIFICATION is not None:
+        try:
+            t = get_swinv2_transforms(256)
+            input_t = t(img.convert('RGB')).unsqueeze(0).to(device)
+            with torch.no_grad():
+                out = SWINV2_CLASSIFICATION(input_t)
+                probs = torch.softmax(out, dim=1)[0].cpu().numpy()
+                top_idx = probs.argmax()
+                if top_idx < len(SWINV2_CATEGORIES):
+                    cat = SWINV2_CATEGORIES[top_idx]
+                    conf = probs[top_idx]
+                    from PIL import Image as _PILImage
+                    # Get HS info if CARGO_DATABASE is available
+                    try:
+                        from __main__ import CARGO_DATABASE, get_hs_info
+                        info = get_hs_info(cat)
+                        hs_txt = f" — {info['hs']} {info['ar']}"
+                    except:
+                        hs_txt = ""
+                    
+                    html += f"<div style='background:#E8EAF6;border:2px solid #3F51B5;border-radius:10px;padding:8px;margin-top:6px;'>"
+                    html += f"<div style='background:#3F51B5;color:white;padding:4px 10px;border-radius:6px;font-size:11px;font-weight:bold;margin-bottom:4px;'>🔱 Swin-V2 Classification — 94.1% F1 (43 classes)</div>"
+                    html += f"<div style='display:flex;align-items:center;gap:8px;'>"
+                    html += f"<div style='background:#3F51B5;color:white;padding:6px 14px;border-radius:8px;font-weight:bold;font-size:14px;'>{cat.upper()}{hs_txt}</div>"
+                    html += f"<div style='flex:1;background:#ddd;border-radius:10px;height:20px;overflow:hidden;'><div style='background:#3F51B5;height:100%;width:{conf*100:.0f}%;border-radius:10px;'></div></div>"
+                    html += f"<div style='font-weight:bold;color:#3F51B5;'>{conf:.0%}</div>"
+                    html += f"</div></div>"
+        except Exception as e:
+            print(f"⚠️ Swin-V2 cls display: {e}")
     
-    if 'CONCEALMENT_MODEL' in dir() and CONCEALMENT_MODEL is not None:
+    # Concealment Detection
+    if SWINV2_CONCEALMENT is not None:
         try:
-            swin_t = get_swin_transforms()
-            img_t = swin_t(img.convert('RGB')).unsqueeze(0).to(DEVICE)
+            conc_size = getattr(SWINV2_CONCEALMENT, '_conc_img_size', 256)
+            t = get_swinv2_transforms(conc_size)
+            input_t = t(img.convert('RGB')).unsqueeze(0).to(device)
             with torch.no_grad():
-                out = CONCEALMENT_MODEL(img_t)
+                out = SWINV2_CONCEALMENT(input_t)
                 probs = torch.softmax(out, dim=1)[0].cpu().numpy()
                 conc_idx = probs.argmax()
                 conc_label = SWINV2_CONCEALMENT_CLASSES[conc_idx]
                 conc_conf = probs[conc_idx]
                 
                 if conc_label == 'no_match':
-                    is_match = False
-                    conc_color = '#E53935'
-                    conc_text = f"⚠️ غير مطابق ({conc_conf:.1%})"
+                    cc = '#E53935'
+                    ct = f"⚠️ غير مطابق ({conc_conf:.1%})"
                 else:
-                    conc_color = '#2E7D32'
-                    conc_text = f"✅ مطابق ({conc_conf:.1%})"
-                    
-                swin_concealment_html = f"<div style='background:{conc_color}15;border:1px solid {conc_color}40;border-radius:8px;padding:6px 10px;margin-top:4px;font-size:12px;'><b style='color:{conc_color};'>🔍 Swin-V2 Concealment (98.9% F1): {conc_text}</b></div>"
-        except:
-            pass
+                    cc = '#2E7D32'
+                    ct = f"✅ مطابق ({conc_conf:.1%})"
+                
+                html += f"<div style='background:{cc}15;border:1px solid {cc}40;border-radius:8px;padding:6px 10px;margin-top:4px;font-size:12px;'>"
+                html += f"<b style='color:{cc};'>🔍 Concealment Detection (98.9% F1): {ct}</b></div>"
+        except Exception as e:
+            print(f"⚠️ Concealment display: {e}")
     
-    if 'RISK_ASSESSMENT_MODEL' in dir() and RISK_ASSESSMENT_MODEL is not None:
+    # Risk Assessment
+    if SWINV2_RISK is not None:
         try:
-            swin_t = get_swin_transforms()
-            img_t = swin_t(img.convert('RGB')).unsqueeze(0).to(DEVICE)
+            t = get_swinv2_transforms(256)
+            input_t = t(img.convert('RGB')).unsqueeze(0).to(device)
             with torch.no_grad():
-                out = RISK_ASSESSMENT_MODEL(img_t)
+                out = SWINV2_RISK(input_t)
                 probs = torch.softmax(out, dim=1)[0].cpu().numpy()
-                risk_idx = probs.argmax()
-                risk_label = SWINV2_RISK_CLASSES[risk_idx]
-                risk_conf = probs[risk_idx]
+                ridx = probs.argmax()
+                rlabel = SWINV2_RISK_CLASSES[ridx]
+                rconf = probs[ridx]
                 
                 risk_ar = {'critical':'حرج','high':'عالي','low':'منخفض','medium':'متوسط','safe':'آمن'}
-                risk_colors = {'critical':'#B71C1C','high':'#D84315','medium':'#F9A825','low':'#558B2F','safe':'#2E7D32'}
-                r_c = risk_colors.get(risk_label, '#666')
+                risk_clrs = {'critical':'#B71C1C','high':'#D84315','medium':'#F9A825','low':'#558B2F','safe':'#2E7D32'}
+                rc2 = risk_clrs.get(rlabel, '#666')
                 
-                swin_risk_html = f"<div style='background:{r_c}15;border:1px solid {r_c}40;border-radius:8px;padding:6px 10px;margin-top:4px;font-size:12px;'><b style='color:{r_c};'>⚡ Swin-V2 Risk (97.2% F1): {risk_ar.get(risk_label, risk_label)} ({risk_conf:.1%})</b></div>"
-        except:
-            pass
-
-
-# ═══════════════════════════════════════════════════════════════
-# التعديل 7: أضف HTML لـ Swin-V2 في قالب HTML
-# ابحث عن: {anomaly_html}
-# أضف بعدها: {swin_concealment_html}
-#            {swin_risk_html}
-# ═══════════════════════════════════════════════════════════════
-
-    # في HTML template، ابحث عن:
-    #   {anomaly_html}
-    # واستبدل بـ:
-    #   {anomaly_html}
-    #   {swin_concealment_html}
-    #   {swin_risk_html}
-
-
-# ═══════════════════════════════════════════════════════════════
-# التعديل 8: أضف Swin-V2 في عنوان التطبيق
-# ابحث عن: 6 Classification + Detection
-# استبدل بـ: 6+1 Classification + 3 Swin-V2 + Detection
-# ═══════════════════════════════════════════════════════════════
-
-# ═══════════════════════════════════════════════════════════════
-# ملخص التعديلات:
-# 1. VERSION = "17.2"
-# 2. أضف SWINV2_CATEGORIES + CONCEALMENT + RISK classes
-# 3. أضف get_swin_transforms()
-# 4. أضف تحميل 3 نماذج Swin-V2 في load_all_models
-# 5. أصلح classify_image لاستخدام transforms مختلفة لـ Swin-V2
-# 6. أضف كشف الإخفاء + تقييم الخطورة في analyze_image
-# 7. أضف عرض النتائج في HTML
-# 8. غيّر العنوان
-#
-# كل التحليل الفيزيائي (8 تقنيات + PRMI + VTA) يبقى كما هو!
-# ═══════════════════════════════════════════════════════════════
+                html += f"<div style='background:{rc2}15;border:1px solid {rc2}40;border-radius:8px;padding:6px 10px;margin-top:4px;font-size:12px;'>"
+                html += f"<b style='color:{rc2};'>⚡ Risk Assessment (97.2% F1): {risk_ar.get(rlabel, rlabel)} ({rconf:.1%})</b></div>"
+        except Exception as e:
+            print(f"⚠️ Risk display: {e}")
+    
+    return html