Upload models.py with huggingface_hub

models.py

@@ -0,0 +1,106 @@
+import timm
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from PIL import Image
+import io
+import os
+
+# --- CONFIGURATION ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+IMG_SIZE = 224
+BATCH_SIZE = 16
+EPOCHS = 1  # Pour test rapide
+MODEL_PATH = "mpox_vit_model_local.pth"
+DATASET_PATH = "Mpox2-1"
+
+# --- TRANSFORMATIONS ---
+transform_train = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(15),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5]*3, [0.5]*3)
+])
+
+transform_val = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5]*3, [0.5]*3)
+])
+
+# --- DATASET LOCAL ---
+train_dataset = datasets.ImageFolder(os.path.join(DATASET_PATH, "train"), transform=transform_train)
+
+if os.path.exists(os.path.join(DATASET_PATH, "valid")):
+    val_dataset = datasets.ImageFolder(os.path.join(DATASET_PATH, "valid"), transform=transform_val)
+else:
+    val_dataset = None
+
+train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE) if val_dataset else None
+
+classes = train_dataset.classes
+num_classes = len(classes)
+
+# --- MODELE ---
+def create_model(pretrained=True):
+    model = timm.create_model('vit_base_patch16_224', pretrained=pretrained)
+    model.head = nn.Linear(model.head.in_features, num_classes)
+    return model
+
+# --- CHARGEMENT DU MODELE ---
+def load_or_train_model():
+    need_train = True
+    if os.path.exists(MODEL_PATH):
+        if os.path.getsize(MODEL_PATH) > 1000:  # Vérifie que le fichier n'est pas vide
+            try:
+                model = create_model(pretrained=False)
+                model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+                model.eval()
+                model.to(device)
+                print(f"Modèle chargé depuis {MODEL_PATH}")
+                need_train = False
+                return model
+            except Exception as e:
+                print(f"Erreur au chargement du modèle: {e}")
+
+    # Entraînement si le modèle n'existe pas ou est corrompu
+    print("Entraînement d'un nouveau modèle...")
+    model = create_model(pretrained=True)
+    model.to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
+    model.train()
+    for epoch in range(EPOCHS):
+        for images, labels in train_loader:
+            images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+    torch.save(model.state_dict(), MODEL_PATH)
+    print(f"Modèle sauvegardé : {MODEL_PATH}")
+    model.eval()
+    return model
+
+# model = load_or_train_model()
+
+# --- PRÉDICTION ---
+def predict_image(image_bytes):
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5]*3, [0.5]*3)
+    ])
+    x = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(x)
+        _, predicted = torch.max(outputs, 1)
+    return classes[predicted.item()]
+
+