LeoFades/COS30082_Herbarium_Field_Classification
0
1"""Plant Species Identification - HuggingFace Space"""
2import os
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6import numpy as np
7from PIL import Image
8from torchvision import transforms
9import gradio as gr
10from huggingface_hub import hf_hub_download
11
12DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
14# ----------------------
15# Transforms
16# ----------------------
17tf_dino = transforms.Compose([
18 transforms.Resize((518, 518)), transforms.CenterCrop(518), transforms.ToTensor(),
19 transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225])
20])
21tf_cnn = transforms.Compose([
22 transforms.Resize((224, 224)), transforms.ToTensor(),
23 transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225])
24])
25
26# ----------------------
27# Model Definitions
28# ----------------------
29class LinearHead(nn.Module):
30 def __init__(self, in_dim=768, num_classes=100):
31 super().__init__()
32 self.fc = nn.Linear(in_dim, num_classes)
33
34 def forward(self, x):
35 return self.fc(x)
36
37
38class TripletProjector(nn.Module):
39 def __init__(self, in_dim=768, proj_dim=1024):
40 super().__init__()
41 self.net = nn.Sequential(
42 nn.Linear(in_dim, 1024),
43 nn.ReLU(),
44 nn.BatchNorm1d(1024),
45 nn.Dropout(0.1),
46 nn.Linear(1024, proj_dim)
47 )
48 def forward(self, x):
49 return F.normalize(self.net(x), dim=1)
50
51class DANNHead(nn.Module):
52 def __init__(self, in_dim=768, proj_dim=512, hid=512, num_classes=100):
53 super().__init__()
54 self.proj = nn.Sequential(
55 nn.Linear(in_dim, hid),
56 nn.ReLU(),
57 nn.Linear(hid, proj_dim)
58 )
59 self.classifier = nn.Linear(proj_dim, num_classes)
60
61 def forward(self, x):
62 z = F.normalize(self.proj(x), dim=-1)
63 logits = self.classifier(z)
64 return z, logits
65
66class PlantCNN(nn.Module):
67 def __init__(self, arch, num_classes=100):
68 super().__init__()
69 if arch == "resnet50":
70 from torchvision.models import resnet50, ResNet50_Weights
71 self.backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
72 in_features = self.backbone.fc.in_features
73 self.backbone.fc = nn.Sequential(
74 nn.Dropout(0.3),
75 nn.Linear(in_features, 256),
76 nn.ReLU(),
77 nn.BatchNorm1d(256),
78 nn.Dropout(0.3),
79 nn.Linear(256, num_classes)
80 )
81
82 elif arch == "efficientnet":
83 from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
84 self.backbone = efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)
85 in_features = self.backbone.classifier[1].in_features
86 self.backbone.classifier = nn.Sequential(
87 nn.Dropout(0.2),
88 nn.Linear(in_features, 256),
89 nn.SiLU(),
90 nn.BatchNorm1d(256),
91 nn.Dropout(0.3),
92 nn.Linear(256, num_classes)
93 )
94
95 else: # mobilenet
96 from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
97 self.backbone = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)
98 in_features = self.backbone.classifier[1].in_features
99 self.backbone.classifier = nn.Sequential(
100 nn.Dropout(0.3),
101 nn.Linear(in_features, 256),
102 nn.ReLU(),
103 nn.BatchNorm1d(256),
104 nn.Dropout(0.3),
105 nn.Linear(256, num_classes)
106 )
107
108 def forward(self, x):
109 return self.backbone(x)
110
111
112# ----------------------
113# Global state
114# ----------------------
115models = {}
116species_map = {}
117label_to_idx = {}
118idx_to_label = {}
119
120# ----------------------
121# Helper functions
122# ----------------------
123def dl(repo, fname):
124 try: return hf_hub_download(repo_id=repo, filename=fname)
125 except: return None
126
127def load_all():
128 global species_map, label_to_idx, idx_to_label
129
130 def valid_path(p):
131 return isinstance(p, str) and os.path.exists(p)
132
133 def safe_dl(repo, fname):
134 try:
135 return dl(repo, fname)
136 except Exception as e:
137 print(f"[dl] failed to download {fname} from {repo}: {e}")
138 return None
139
140 # --- Load metadata ---
141 base_repo = os.environ.get("BASE_REPO", "")
142 if base_repo:
143 # species list
144 p = safe_dl(base_repo, "species_list.txt")
145 if valid_path(p):
146 with open(p, "r", encoding="utf-8") as f:
147 for line in f:
148 parts = line.strip().split("; ")
149 if len(parts) >= 2:
150 try: sid = int(parts[0])
151 except: continue
152 species_map[sid] = "; ".join(parts[1:])
153 # label mapping
154 p = safe_dl(base_repo, "label_to_idx.pt")
155 if valid_path(p):
156 data = torch.load(p, map_location="cpu")
157 if isinstance(data, dict):
158 label_to_idx.update(data)
159 idx_to_label.update({v:k for k,v in data.items()})
160
161 # --- DINOv2 ---
162 try:
163 import timm
164 models["dinov2"] = timm.create_model("vit_base_patch14_reg4_dinov2", pretrained=True)
165 repo = os.environ.get("DINOV2_REPO", "")
166 if repo:
167 p = safe_dl(repo, "dinov2_backbone.pth")
168 if valid_path(p):
169 models["dinov2"].load_state_dict(torch.load(p, map_location="cpu"), strict=False)
170 models["dinov2"].eval().to(DEVICE)
171 for param in models["dinov2"].parameters(): param.requires_grad = False
172 except Exception as e:
173 print("[load_all] DINOv2 load failed:", e)
174 models.pop("dinov2", None)
175
176 # --- Linear ---
177 repo = os.environ.get("LINEAR_REPO", "")
178 if repo:
179 p = safe_dl(repo, "linear_classifier.pth")
180 if valid_path(p):
181 models["linear"] = LinearHead(768, len(label_to_idx) or 100)
182 ckpt = torch.load(p, map_location="cpu", weights_only=False)
183 if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
184 models["linear"].load_state_dict(ckpt["model_state_dict"])
185 else:
186 models["linear"].load_state_dict(ckpt)
187 models["linear"].eval().to(DEVICE)
188
189 # --- Triplet ---
190 repo = os.environ.get("TRIPLET_REPO", "")
191 if repo:
192 p = safe_dl(repo, "best_projector.pth")
193 if valid_path(p):
194 models["triplet"] = TripletProjector(768, 1024)
195 models["triplet"].load_state_dict(torch.load(p, map_location="cpu"))
196 models["triplet"].eval().to(DEVICE)
197 p = safe_dl(repo, "herbarium_embeddings.pt")
198 if valid_path(p):
199 data = torch.load(p, map_location="cpu")
200 if isinstance(data, dict):
201 # trained as dict
202 emb = data.get("embeddings", data.get("emb"))
203 ids = data.get("species_ids", data.get("labels"))
204 elif isinstance(data, torch.Tensor):
205 # trained directly as tensor
206 emb = data
207 p2 = safe_dl(repo, "herbarium_species_ids.pt")
208 if valid_path(p2):
209 ids = torch.load(p2, map_location="cpu")
210 else:
211 ids = None
212 else:
213 emb = None
214 ids = None
215
216 if emb is not None:
217 models["herb_emb"] = emb if isinstance(emb, torch.Tensor) else torch.tensor(emb)
218 if ids is not None:
219 models["herb_ids"] = torch.tensor(ids) if not isinstance(ids, torch.Tensor) else ids
220
221
222 # --- DANN ---
223 repo = os.environ.get("DANN_REPO", "")
224 if repo:
225 p = safe_dl(repo, "best_head_dann.pth")
226 if valid_path(p):
227 ckpt = torch.load(p, map_location="cpu", weights_only=False)
228 if isinstance(ckpt, dict):
229 models["dann"] = DANNHead(768, 512, 512, len(label_to_idx) or 100)
230 head_state = ckpt.get("head_state", ckpt)
231 models["dann"].load_state_dict(head_state)
232 models["dann"].eval().to(DEVICE)
233 if "prototypes" in ckpt:
234 models["prototypes"] = ckpt["prototypes"]
235
236 # --- CNNs ---
237 for name, env_var, arch in [("resnet50","RESNET_REPO","resnet50"),
238 ("efficientnet","EFFICIENTNET_REPO","efficientnet"),
239 ("mobilenetv2","MOBILENET_REPO","mobilenetv2")]:
240 repo = os.environ.get(env_var,"")
241 if not repo: continue
242 fname = f"{arch}_plant.pth" if "efficient" in arch else f"{name}_plant.pth"
243 p = safe_dl(repo,fname)
244 if valid_path(p):
245 models[name] = PlantCNN(arch,len(label_to_idx) or 100)
246 models[name].load_state_dict(torch.load(p,map_location="cpu"))
247 models[name].eval().to(DEVICE)
248
249# ----------------------
250# Feature extraction
251# ----------------------
252@torch.no_grad()
253def get_dino_feat(img):
254 if "dinov2" not in models: return None
255 x = tf_dino(img).unsqueeze(0).to(DEVICE)
256 return models["dinov2"](x)
257
258# ----------------------
259# Prediction formatting
260# ----------------------
261def format_result(method, preds):
262 if preds is None: return f"**{method}**: Not available\n\n"
263 lines = [f"**{method}**"]
264 for i,(idx,score) in enumerate(preds[:5]):
265 sid = idx_to_label.get(int(idx), int(idx)) if idx_to_label else int(idx)
266 name = species_map.get(sid,f"Species {sid}")
267 lines.append(f"{i+1}. {name} ({float(score)*100:.1f}%)")
268 return "\n".join(lines)+"\n\n"
269
270# ----------------------
271# Prediction function
272# ----------------------
273@torch.no_grad()
274def predict(image, methods):
275 if image is None: return "Upload an image"
276 if isinstance(image, np.ndarray): image = Image.fromarray(image).convert("RGB")
277
278 results = ""
279 feat = get_dino_feat(image)
280
281 for m in methods:
282 preds = None
283 try:
284 if m=="Linear" and "linear" in models and feat is not None:
285 logits = models["linear"](feat)
286 probs = F.softmax(logits, dim=1)
287 vals, idxs = probs.topk(5)
288 preds = list(zip(idxs[0].cpu().numpy(), vals[0].cpu().numpy()))
289
290 elif m=="Triplet" and "triplet" in models and "herb_emb" in models and feat is not None:
291 q = models["triplet"](feat)
292 h = models["triplet"](models["herb_emb"].to(DEVICE))
293 sim = (q @ h.t())[0]
294 vals, idxs = sim.topk(5)
295 preds = [(models["herb_ids"][i].item(), (v.item()+1)/2) for i,v in zip(idxs, vals)]
296
297 elif m=="DANN+Proto" and "dann" in models and feat is not None:
298 z, logits = models["dann"](feat)
299 if "prototypes" in models:
300 probs = F.softmax(z @ models["prototypes"].to(DEVICE).t() * 10, dim=1)
301 else:
302 probs = F.softmax(logits, dim=1)
303 vals, idxs = probs.topk(5)
304 preds = list(zip(idxs[0].cpu().numpy(), vals[0].cpu().numpy()))
305
306 elif m=="ResNet50" and "resnet" in models:
307 x = tf_cnn(image).unsqueeze(0).to(DEVICE)
308 probs = F.softmax(models["resnet"](x), dim=1)
309 vals, idxs = probs.topk(5)
310 preds = list(zip(idxs[0].cpu().numpy(), vals[0].cpu().numpy()))
311
312 elif m=="EfficientNet" and "efficientnet" in models:
313 x = tf_cnn(image).unsqueeze(0).to(DEVICE)
314 probs = F.softmax(models["efficientnet"](x), dim=1)
315 vals, idxs = probs.topk(5)
316 preds = list(zip(idxs[0].cpu().numpy(), vals[0].cpu().numpy()))
317
318 elif m=="MobileNetV2" and "mobilenet" in models:
319 x = tf_cnn(image).unsqueeze(0).to(DEVICE)
320 probs = F.softmax(models["mobilenet"](x), dim=1)
321 vals, idxs = probs.topk(5)
322 preds = list(zip(idxs[0].cpu().numpy(), vals[0].cpu().numpy()))
323 except Exception as e:
324 preds = None
325 results += format_result(m,preds)
326
327 return results if results else "Select at least one method"
328
329
330# ----------------------
331# Make sure models are loaded first
332# ----------------------
333load_all()
334
335# ----------------------
336# Gradio interface
337# ----------------------
338with gr.Blocks(title="🌱 Plant Species Identification") as demo:
339
340 gr.Markdown("## Cross-domain plant species identification using herbarium-field matching")
341
342 with gr.Row():
343 img_input = gr.Image(type="pil", label="Upload Plant Image")
344 methods_input = gr.CheckboxGroup(
345 ["ResNet50", "EfficientNet", "MobileNetV2", "Linear", "Triplet", "DANN+Proto"],
346 value=["Linear", "Triplet", "DANN+Proto"], # default selected
347 label="Methods"
348 )
349
350 output_md = gr.Markdown(label="Results")
351
352 submit_btn = gr.Button("Predict")
353 submit_btn.click(
354 fn=predict,
355 inputs=[img_input, methods_input],
356 outputs=output_md
357 )
358
359if __name__ == "__main__":
360 demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
361
362
363