Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -textimages/BLAST1_008.jpg filter=lfs diff=lfs merge=lfs -text
+images/BLAST1_016.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,11 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+# checkpoints

.python-version

@@ -0,0 +1 @@
+3.12

README.md

@@ -1,3 +1,66 @@
----
-license: mit
----
+---
+language: en
+tags:
+- pytorch
+- computer-vision
+- image-classification
+- rice-disease
+license: mit
+---
+
+# VGG16-CNN Rice Disease Classification Model
+
+This model is designed for classifying rice plant diseases using a modified VGG16 architecture with additional CNN layers.
+
+## Model Description
+
+### Architecture
+- Base model: VGG16 (pretrained on ImageNet)
+- Additional custom CNN layer with:
+  - Conv2d(512, 64, kernel_size=3)
+  - ReLU activation
+  - BatchNorm2d
+  - MaxPool2d
+- Custom classifier with:
+  - Linear layers (32*3*6 → 1024 → 5)
+  - Dropout (0.4)
+
+### Task
+Image classification for rice plant diseases
+
+### Classes
+1. Bacterialblight
+2. Blast
+3. Brownspot
+4. Healthy
+5. Tungro
+
+## Training
+
+The model uses transfer learning with a frozen VGG16 backbone.
+
+## Intended Use
+- Primary intended use: Rice disease diagnosis through leaf image analysis
+- Out-of-scope use: Should not be used for critical agricultural decisions without expert verification
+
+## Input
+- RGB images
+- Required size: 224x224 pixels
+- Preprocessing:
+  - Normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+## Limitations
+Please note that this model should be used as a supportive tool and not as a sole decision-maker for disease diagnosis.
+
+## Model Author
+[Your Name/Organization]
+
+## Citation
+If you use this model, please cite:
+```
+@software{vgg_cnn_rice_disease,
+  title={VGG16-CNN Rice Disease Classification Model},
+  version={0.1.0},
+  year={2024}
+}
+```

checkpoints/vgg_net_model_50.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9da3747b37fa8286a07b7dd7eb5cfdfd734e46fcb94b099c346749625e841e05
+size 557011773

main.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+from PIL import Image
+from torch import Tensor
+from torchvision import transforms
+
+from model import VGG16WithCNN
+
+
+def getModel(device: torch.device, model_path: str):
+    model = VGG16WithCNN(5)
+    # 加载训练好的权重
+    model.load_state_dict(
+        torch.load(
+            model_path,
+            weights_only=True,
+        )
+    )
+
+    model.to(device)
+    return model
+
+
+def preprocess_image(image_path: str, image_size=(224, 224)):
+    """
+    预处理图片，使其符合模型输入要求
+    """
+    transform = transforms.Compose(
+        [
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    # 打开图片并转换
+    image: Image.Image = Image.open(image_path).convert("RGB")
+    image_tensor: Tensor = transform(image)
+    # 添加batch维度
+    image_tensor = image_tensor.unsqueeze(0)
+
+    return image_tensor
+
+
+def predict_single_image(
+    image_path: str, model: nn.Module, device: torch.device, class_names: list[str]
+) -> str:
+    """
+    预测单个图片的标签
+    Args:
+        image_path: 图片路径
+        model: 模型
+        device: 设备
+
+    Returns:
+        预测的标签名
+    """
+
+    image_tensor = preprocess_image(image_path)
+
+    image_tensor = image_tensor.to(device)
+    # 预测
+    model.eval()
+    with torch.no_grad():
+        output = model(image_tensor)
+        _, pred = torch.max(output, 1)
+
+    predicted_label = class_names[int(pred.item())]
+
+    return predicted_label
+
+
+if __name__ == "__main__":
+    # 测试单张图片预测
+    # 注意：需要替换为实际的测试图片路径
+
+    p = "./checkpoints/vgg_net_model_50.pth"
+
+    class_names = [
+        "Bacterialblight",
+        "Blast",
+        "Brownspot",
+        "Healthy",
+        "Tungro",
+    ]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = getModel(device=device, model_path=p)
+
+    test_image_path = "./images/BLAST1_011.jpg"
+    try:
+        predicted_label = predict_single_image(
+            test_image_path, model, device, class_names=class_names
+        )
+        print("\nSingle image prediction result:")
+        print(f"Image: {test_image_path}")
+        print(f"Predicted label: {predicted_label}")
+    except FileNotFoundError:
+        print("Please provide a valid image path to test single image prediction")

model.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from PIL import Image
+from torch import Tensor, nn
+from torchvision.models import VGG, VGG16_Weights
+
+
+class VGG16WithCNN(nn.Module):
+    def __init__(self, num_classes: int = 10):
+        super(VGG16WithCNN, self).__init__()
+        self.num_classes = num_classes
+        self.vgg16: VGG = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
+
+        for param in self.vgg16.parameters():
+            param.requires_grad = False
+
+        self.custom_cnn = nn.Sequential(
+            nn.Conv2d(512, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(64),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+
+        # 分类器
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(32 * 3 * 6, 1024),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(1024, num_classes),
+        )
+
+    def forward(self, x: Tensor):
+        x = self.vgg16.features(x)
+        x = self.custom_cnn(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x

pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "vgg-cnn"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "huggingface-hub[cli]>=0.35.3",
+    "torch>=2.8.0",
+    "torchvision>=0.23.0",
+]