Create README.md

README.md

@@ -0,0 +1,223 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+import timm
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from tqdm.auto import tqdm
+import wandb
+
+class PlantDiseaseDataset(Dataset):
+    def __init__(self, image_paths, labels, transform=None):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.transform = transform
+        
+    def __len__(self):
+        return len(self.image_paths)
+    
+    def __getitem__(self, idx):
+        image_path = self.image_paths[idx]
+        image = Image.open(image_path).convert('RGB')
+        label = self.labels[idx]
+        
+        if self.transform:
+            image = self.transform(image)
+            
+        return image, label
+
+class PlantDiseaseClassifier:
+    def __init__(self, data_dir, model_name='vit_base_patch16_224', num_classes=38):
+        self.data_dir = Path(data_dir)
+        self.model_name = model_name
+        self.num_classes = num_classes
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Initialize wandb
+        wandb.init(project="plant-disease-classification")
+        
+    def prepare_data(self):
+        """Prepare dataset and create data loaders"""
+        # Data augmentation and normalization for training
+        train_transform = transforms.Compose([
+            transforms.RandomResizedCrop(224),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomVerticalFlip(),
+            transforms.RandomRotation(20),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+        
+        # Just normalization for validation/testing
+        val_transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+        
+        # Collect all image paths and labels
+        image_paths = []
+        labels = []
+        self.class_to_idx = {}
+        
+        for idx, class_dir in enumerate(sorted(self.data_dir.glob('*'))):
+            if class_dir.is_dir():
+                self.class_to_idx[class_dir.name] = idx
+                for img_path in class_dir.glob('*.jpg'):
+                    image_paths.append(str(img_path))
+                    labels.append(idx)
+        
+        # Split data
+        train_paths, val_paths, train_labels, val_labels = train_test_split(
+            image_paths, labels, test_size=0.2, stratify=labels, random_state=42
+        )
+        
+        # Create datasets
+        train_dataset = PlantDiseaseDataset(train_paths, train_labels, train_transform)
+        val_dataset = PlantDiseaseDataset(val_paths, val_labels, val_transform)
+        
+        # Create data loaders
+        self.train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
+        self.val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
+        
+        return self.train_loader, self.val_loader
+    
+    def create_model(self):
+        """Initialize the Vision Transformer model"""
+        self.model = timm.create_model(
+            self.model_name,
+            pretrained=True,
+            num_classes=self.num_classes
+        )
+        self.model = self.model.to(self.device)
+        
+        # Loss function and optimizer
+        self.criterion = nn.CrossEntropyLoss()
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=2e-5,
+            weight_decay=0.01
+        )
+        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            self.optimizer,
+            T_max=10
+        )
+        
+        return self.model
+    
+    def train_epoch(self, epoch):
+        """Train for one epoch"""
+        self.model.train()
+        total_loss = 0
+        correct = 0
+        total = 0
+        
+        progress_bar = tqdm(self.train_loader, desc=f'Epoch {epoch}')
+        
+        for batch_idx, (inputs, targets) in enumerate(progress_bar):
+            inputs, targets = inputs.to(self.device), targets.to(self.device)
+            
+            self.optimizer.zero_grad()
+            outputs = self.model(inputs)
+            loss = self.criterion(outputs, targets)
+            
+            loss.backward()
+            self.optimizer.step()
+            
+            total_loss += loss.item()
+            _, predicted = outputs.max(1)
+            total += targets.size(0)
+            correct += predicted.eq(targets).sum().item()
+            
+            progress_bar.set_postfix({
+                'Loss': total_loss/(batch_idx+1),
+                'Acc': 100.*correct/total
+            })
+            
+            # Log to wandb
+            wandb.log({
+                'train_loss': loss.item(),
+                'train_acc': 100.*correct/total
+            })
+            
+        return total_loss/len(self.train_loader), 100.*correct/total
+    
+    def validate(self):
+        """Validate the model"""
+        self.model.eval()
+        total_loss = 0
+        correct = 0
+        total = 0
+        
+        with torch.no_grad():
+            for inputs, targets in tqdm(self.val_loader, desc='Validating'):
+                inputs, targets = inputs.to(self.device), targets.to(self.device)
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, targets)
+                
+                total_loss += loss.item()
+                _, predicted = outputs.max(1)
+                total += targets.size(0)
+                correct += predicted.eq(targets).sum().item()
+                
+        accuracy = 100.*correct/total
+        avg_loss = total_loss/len(self.val_loader)
+        
+        # Log to wandb
+        wandb.log({
+            'val_loss': avg_loss,
+            'val_acc': accuracy
+        })
+        
+        return avg_loss, accuracy
+    
+    def train(self, epochs=10):
+        """Complete training process"""
+        best_acc = 0
+        
+        for epoch in range(epochs):
+            train_loss, train_acc = self.train_epoch(epoch)
+            val_loss, val_acc = self.validate()
+            self.scheduler.step()
+            
+            print(f'\nEpoch {epoch}:')
+            print(f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%')
+            print(f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%')
+            
+            # Save best model
+            if val_acc > best_acc:
+                best_acc = val_acc
+                torch.save({
+                    'model_state_dict': self.model.state_dict(),
+                    'optimizer_state_dict': self.optimizer.state_dict(),
+                    'class_to_idx': self.class_to_idx
+                }, 'best_model.pth')
+        
+        wandb.finish()
+    
+    def save_for_huggingface(self):
+        """Save model in HF中国镜像站 format"""
+        # Load best model
+        checkpoint = torch.load('best_model.pth')
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        
+        # Save model and config
+        self.model.save_pretrained('plant_disease_model')
+        
+        # Save class mapping
+        idx_to_class = {v: k for k, v in self.class_to_idx.items()}
+        pd.Series(idx_to_class).to_json('class_mapping.json')
+
+if __name__ == "__main__":
+    classifier = PlantDiseaseClassifier(data_dir="path/to/dataset")
+    classifier.prepare_data()
+    classifier.create_model()
+    classifier.train(epochs=10)
+    classifier.save_for_huggingface()