Hugging Face's logo

chaosbee997
/
rice-seed-classifier

ml-intern
Model card Files
xet
 Community
rice-seed-classifier
File size: 5,045 Bytes
2807a2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
"""
Colab-ready training script for rice seed/grain classification.

Quick start in Google Colab:
  1. Go to https://colab.research.google.com
  2. Create a new notebook
  3. Copy each "cell" below into a Colab code cell and run sequentially.

Hardware: Runtime → Change runtime type → T4 GPU (recommended)
Expected time: ~15-30 minutes for 5 epochs on T4.
"""

# ===================== CELL 1: 检查GPU =====================
# !nvidia-smi

# ===================== CELL 2: 安装依赖 =====================
# !pip install -q transformers datasets accelerate evaluate pillow huggingface_hub
# import os
# os.kill(os.getpid(), 9)  # 安装后重启runtime（可选，transformers建议）

# ===================== CELL 3: 导入库 =====================
import os
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
)
from PIL import Image
import evaluate
from huggingface_hub import notebook_login

# ===================== CELL 4: 登录HF (需要Write权限Token) =====================
# notebook_login()  # 交互式弹窗输入token

# ===================== CELL 5: 加载数据集 =====================
DATASET_NAME = "nateraw/rice-image-dataset"
print(f"Loading dataset: {DATASET_NAME}...")
ds = load_dataset(DATASET_NAME)
print(ds)

labels = ds["train"].features["label"].names
num_labels = len(labels)
print(f"Classes ({num_labels}): {labels}")

# ===================== CELL 6: 划分训练/验证集 =====================
split_ds = ds["train"].train_test_split(test_size=0.15, stratify_by_column="label", seed=42)
train_ds = split_ds["train"]
val_ds = split_ds["test"]
print(f"Train: {len(train_ds)} | Val: {len(val_ds)}")

# ===================== CELL 7: 加载模型 =====================
MODEL_NAME = "microsoft/resnet-18"
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

label2id = {label: i for i, label in enumerate(labels)}
id2label = {i: label for i, label in enumerate(labels)}

model = AutoModelForImageClassification.from_pretrained(
    MODEL_NAME,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True,
)
print(f"Model loaded: {MODEL_NAME} ({sum(p.numel() for p in model.parameters())/1e6:.1f}M params)")

# ===================== CELL 8: 预处理函数 =====================
def transform(example_batch):
    images = []
    for img in example_batch["image"]:
        if isinstance(img, Image.Image):
            if img.mode != "RGB":
                img = img.convert("RGB")
            images.append(img)
        else:
            img = Image.fromarray(np.array(img)).convert("RGB")
            images.append(img)
    inputs = processor(images, return_tensors="pt")
    inputs["labels"] = example_batch["label"]
    return inputs

train_ds.set_transform(transform)
val_ds.set_transform(transform)

# ===================== CELL 9: 评估指标 =====================
accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    preds = np.argmax(predictions, axis=1)
    acc = accuracy.compute(predictions=preds, references=labels)
    f1_score = f1.compute(predictions=preds, references=labels, average="weighted")
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

# ===================== CELL 10: 训练参数 =====================
OUTPUT_REPO = "chaosbee997/rice-seed-classifier"  # 改成你的仓库名

args = TrainingArguments(
    output_dir="/content/rice-seed-classifier",
    remove_unused_columns=False,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    num_train_epochs=5,
    warmup_ratio=0.1,
    logging_strategy="steps",
    logging_steps=50,
    logging_first_step=True,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    seed=42,
    push_to_hub=True,
    hub_model_id=OUTPUT_REPO,
    report_to="none",
)

# ===================== CELL 11: 训练 =====================
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    compute_metrics=compute_metrics,
    data_collator=DefaultDataCollator(),
    tokenizer=processor,
)

print("Starting training...")
trainer.train()

# ===================== CELL 12: 评估 + 上传 =====================
metrics = trainer.evaluate()
print("Evaluation:", metrics)
trainer.push_to_hub()
print(f"Model uploaded to: https://huggingface.co/{OUTPUT_REPO}")

# ===================== CELL 13: 推理测试 =====================
from transformers import pipeline
import matplotlib.pyplot as plt

classifier = pipeline("image-classification", model=OUTPUT_REPO)
test_image = ds["train"][100]["image"]

plt.figure(figsize=(5,5))
plt.imshow(test_image)
plt.axis("off")
plt.show()

results = classifier(test_image)
for r in results:
    print(f"  {r['label']}: {r['score']*100:.2f}%")