Add custom trained hybrid model for Hen and Peacock (using LFS)

.gitattributes

@@ -0,0 +1 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -4,3 +4,4 @@ uploads/
 .vscode/
 *.pyc
 BirdNET-Analyzer-main*/
+dataset/

app.py

@@ -39,6 +39,9 @@ os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
 models = {
     'image_processor': None,
     'image_model': None,
+    'custom_processor': None,
+    'custom_model': None,
+    'custom_loaded': False,
     'loaded': False,
     'loading': False,
     'error': None
@@ -60,9 +63,23 @@ def load_models():
 
         models['image_processor'] = AutoImageProcessor.from_pretrained("chriamue/bird-species-classifier")
         models['image_model'] = AutoModelForImageClassification.from_pretrained("chriamue/bird-species-classifier")
+        
+        # Try loading local custom fine-tuned model if it exists
+        custom_dir = "custom_model"
+        if os.path.exists(custom_dir) and os.path.exists(os.path.join(custom_dir, "config.json")):
+            print("Found custom model locally. Loading...")
+            try:
+                models['custom_processor'] = AutoImageProcessor.from_pretrained(custom_dir)
+                models['custom_model'] = AutoModelForImageClassification.from_pretrained(custom_dir)
+                models['custom_loaded'] = True
+                print("Custom model loaded successfully.")
+            except Exception as custom_err:
+                print(f"Error loading custom model: {custom_err}")
+                models['custom_loaded'] = False
+
         models['loaded'] = True
         models['error'] = None
-        print("Image model loaded successfully.")
+        print("Image models loaded successfully.")
     except Exception as e:
         models['loaded'] = False
         models['error'] = str(e)
@@ -250,6 +267,27 @@ def classify_image():
         file.save(filepath)
 
         image = Image.open(filepath).convert("RGB")
+
+        # 1. Attempt classification with local custom model first (if loaded)
+        if models['custom_loaded'] and models['custom_model'] is not None:
+            custom_inputs = models['custom_processor'](images=image, return_tensors="pt")
+            with torch.no_grad():
+                custom_outputs = models['custom_model'](**custom_inputs)
+            
+            custom_logits = custom_outputs.logits
+            probs = torch.softmax(custom_logits, dim=-1)
+            pred_idx = torch.argmax(probs, dim=-1).item()
+            label = models['custom_model'].config.id2label[pred_idx]
+            confidence = probs[0][pred_idx].item()
+
+            print(f"Custom model prediction: {label} (confidence: {confidence:.4f})")
+
+            # Route to custom predictions if it's HEN or PEACOCK and confidence is high
+            if label in ["HEN", "PEACOCK"] and confidence >= 0.70:
+                species_display = "Hen" if label == "HEN" else "Peacock"
+                return jsonify({'species': species_display, 'type': 'image'})
+
+        # 2. Fallback to original online model
         inputs = models['image_processor'](images=image, return_tensors="pt")
 
         with torch.no_grad():

custom_model/config.json

@@ -0,0 +1,93 @@
+{
+  "architectures": [
+    "EfficientNetForImageClassification"
+  ],
+  "batch_norm_eps": 0.001,
+  "batch_norm_momentum": 0.99,
+  "depth_coefficient": 1.2,
+  "depth_divisor": 8,
+  "depthwise_padding": [
+    5,
+    8,
+    16
+  ],
+  "drop_connect_rate": 0.2,
+  "dropout_rate": 0.3,
+  "dtype": "float32",
+  "expand_ratios": [
+    1,
+    6,
+    6,
+    6,
+    6,
+    6,
+    6
+  ],
+  "hidden_act": "swish",
+  "hidden_dim": 1408,
+  "id2label": {
+    "0": "HEN",
+    "1": "OTHER",
+    "2": "PEACOCK"
+  },
+  "image_size": 260,
+  "in_channels": [
+    32,
+    16,
+    24,
+    40,
+    80,
+    112,
+    192
+  ],
+  "initializer_range": 0.02,
+  "kernel_sizes": [
+    3,
+    3,
+    5,
+    3,
+    5,
+    5,
+    3
+  ],
+  "label2id": {
+    "HEN": 0,
+    "OTHER": 1,
+    "PEACOCK": 2
+  },
+  "model_type": "efficientnet",
+  "num_block_repeats": [
+    1,
+    2,
+    2,
+    3,
+    3,
+    4,
+    1
+  ],
+  "num_channels": 3,
+  "num_hidden_layers": 64,
+  "out_channels": [
+    16,
+    24,
+    40,
+    80,
+    112,
+    192,
+    320
+  ],
+  "pooling_type": "mean",
+  "problem_type": "single_label_classification",
+  "squeeze_expansion_ratio": 0.25,
+  "strides": [
+    1,
+    2,
+    2,
+    2,
+    1,
+    2,
+    1
+  ],
+  "transformers_version": "4.57.1",
+  "width_coefficient": 1.1
+}

custom_model/model.safetensors

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

custom_model/preprocessor_config.json

@@ -0,0 +1,29 @@
+{
+  "crop_size": {
+    "height": 289,
+    "width": 289
+  },
+  "do_center_crop": false,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "EfficientNetImageProcessor",
+  "image_std": [
+    0.47853944,
+    0.4732864,
+    0.47434163
+  ],
+  "include_top": true,
+  "resample": 0,
+  "rescale_factor": 0.00392156862745098,
+  "rescale_offset": false,
+  "size": {
+    "height": 260,
+    "width": 260
+  }
+}

main.py

@@ -6,6 +6,7 @@ from transformers import AutoImageProcessor, AutoModelForImageClassification
 from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 import librosa
 import threading
+import os
 
 class BirdClassifierApp:
     def __init__(self, root):
@@ -16,6 +17,8 @@ class BirdClassifierApp:
         
         self.image_model = None
         self.image_processor = None
+        self.custom_model = None
+        self.custom_processor = None
         self.audio_model = None
         self.audio_extractor = None
         
@@ -118,6 +121,17 @@ class BirdClassifierApp:
                 self.image_processor = AutoImageProcessor.from_pretrained("chriamue/bird-species-classifier")
                 self.image_model = AutoModelForImageClassification.from_pretrained("chriamue/bird-species-classifier")
                 
+                # Load custom fine-tuned model if available
+                custom_dir = "custom_model"
+                if os.path.exists(custom_dir) and os.path.exists(os.path.join(custom_dir, "config.json")):
+                    try:
+                        self.custom_processor = AutoImageProcessor.from_pretrained(custom_dir)
+                        self.custom_model = AutoModelForImageClassification.from_pretrained(custom_dir)
+                        print("Custom local model loaded.")
+                    except Exception as custom_err:
+                        print(f"Error loading custom local model: {custom_err}")
+                        self.custom_model = None
+
                 self.audio_extractor = AutoFeatureExtractor.from_pretrained("greenarcade/wav2vec2-vd-bird-sound-classification")
                 self.audio_model = AutoModelForAudioClassification.from_pretrained("greenarcade/wav2vec2-vd-bird-sound-classification")
                 
@@ -164,14 +178,34 @@ class BirdClassifierApp:
             self.preview_label.config(image=photo, text="")
             self.preview_label.image = photo
             
-            inputs = self.image_processor(images=image, return_tensors="pt")
-            
-            with torch.no_grad():
-                outputs = self.image_model(**inputs)
+            label = None
             
-            logits = outputs.logits
-            pred = torch.argmax(logits, dim=-1).item()
-            label = self.image_model.config.id2label[pred]
+            # 1. Attempt classification with custom model first
+            if self.custom_model is not None:
+                custom_inputs = self.custom_processor(images=image, return_tensors="pt")
+                with torch.no_grad():
+                    custom_outputs = self.custom_model(**custom_inputs)
+                
+                custom_logits = custom_outputs.logits
+                probs = torch.softmax(custom_logits, dim=-1)
+                pred_idx = torch.argmax(probs, dim=-1).item()
+                custom_label = self.custom_model.config.id2label[pred_idx]
+                confidence = probs[0][pred_idx].item()
+                
+                print(f"Custom model prediction: {custom_label} (confidence: {confidence:.4f})")
+                
+                if custom_label in ["HEN", "PEACOCK"] and confidence >= 0.70:
+                    label = "Hen" if custom_label == "HEN" else "Peacock"
+
+            # 2. Fallback to original online model
+            if label is None:
+                inputs = self.image_processor(images=image, return_tensors="pt")
+                with torch.no_grad():
+                    outputs = self.image_model(**inputs)
+                
+                logits = outputs.logits
+                pred = torch.argmax(logits, dim=-1).item()
+                label = self.image_model.config.id2label[pred]
             
             self.result_label.config(text=f"✅ Predicted Species: {label}")
             self.status_label.config(text="✅ Classification complete!")

train.py

@@ -0,0 +1,235 @@
+import argparse
+import os
+import sys
+import urllib.request
+from pathlib import Path
+from PIL import Image
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# List of public domain images for the 'OTHER' class (using Unsplash to avoid rate limits)
+OTHER_BIRD_URLS = [
+    f"https://images.unsplash.com/{photo_id}?w=500&auto=format&fit=crop&q=80"
+    for photo_id in [
+        "photo-1452570053594-1b985d6ea890",
+        "photo-1480044965905-02098d419e96",
+        "photo-1516233758813-a38d024919c5",
+        "photo-1551085254-e96b210db58a",
+        "photo-1522441815192-d9f04eb0615c",
+        "photo-1506220926022-cc5c12abdb35",
+        "photo-1518998053901-5348d3961a04",
+        "photo-1511823794984-b87716139b88",
+        "photo-1470116890351-be0a9b418409",
+        "photo-1539664030485-a936c7d29fc0",
+        "photo-1444464666168-49d633b86797",
+        "photo-1504386106331-3e4e71712b38",
+        "photo-1555041469-a586c61ea9bc",
+        "photo-1525462519782-b55cef2f882a",
+        "photo-1509023467868-1c40786379f6",
+        "photo-1454496522488-7a8e488e8606",
+        "photo-1549488344-1f9b8d2bd1f3",
+        "photo-1528183429752-a97d0bf99b5a",
+        "photo-1510137600163-2729bc695ac1",
+        "photo-1465153690352-10c1b295ec7e",
+        "photo-1497250681960-ef046c08a56e",
+        "photo-1526336024438-db9b601fc211",
+        "photo-1534067783941-51c9c23eccfd",
+        "photo-1560015534-eca11e59876a",
+        "photo-1551806235-a05ff789c3c7",
+    ]
+]
+
+def download_image(url, filepath):
+    """Downloads an image from a URL with a standard browser User-Agent header."""
+    req = urllib.request.Request(
+        url,
+        headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'}
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=15) as response:
+            with open(filepath, 'wb') as f:
+                f.write(response.read())
+        # Verify it is a valid image
+        with Image.open(filepath) as img:
+            img.verify()
+        return True
+    except Exception as e:
+        if os.path.exists(filepath):
+            os.remove(filepath)
+        print(f"  [Warning] Failed to download or verify {url}: {e}")
+        return False
+
+def setup_other_class(other_dir):
+    """Sets up the OTHER class directory and downloads sample bird images."""
+    os.makedirs(other_dir, exist_ok=True)
+    existing_files = [f for f in os.listdir(other_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    if len(existing_files) >= 10:
+        print(f"-> 'OTHER' directory already has {len(existing_files)} images. Skipping download.")
+        return
+
+    print("-> Downloading background bird images for 'OTHER' class (to avoid forgetting)...")
+    downloaded_count = 0
+    for i, url in enumerate(OTHER_BIRD_URLS):
+        dest_path = os.path.join(other_dir, f"other_bird_{i}.jpg")
+        print(f"  Downloading image {i+1}/{len(OTHER_BIRD_URLS)}...")
+        if download_image(url, dest_path):
+            downloaded_count += 1
+            
+    print(f"-> Completed downloading. Added {downloaded_count} images to 'OTHER' directory.")
+
+class BirdDataset(Dataset):
+    def __init__(self, image_paths, labels, image_processor):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.image_processor = image_processor
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.image_paths[idx]
+        label = self.labels[idx]
+        try:
+            image = Image.open(img_path).convert("RGB")
+            inputs = self.image_processor(images=image, return_tensors="pt")
+            pixel_values = inputs["pixel_values"].squeeze(0)
+            return pixel_values, torch.tensor(label, dtype=torch.long)
+        except Exception as e:
+            # Fallback for corrupt images during training
+            # We return a dummy image (zeros) and the label
+            print(f"  [Warning] Error loading image {img_path}: {e}. Using zero-filled tensor.")
+            dummy_pixel = torch.zeros((3, 224, 224))
+            return dummy_pixel, torch.tensor(label, dtype=torch.long)
+
+def train_model(args):
+    dataset_path = Path(args.dataset_dir)
+    hen_dir = dataset_path / "HEN"
+    peacock_dir = dataset_path / "PEACOCK"
+    other_dir = dataset_path / "OTHER"
+
+    # 1. Verify dataset structure
+    if not hen_dir.exists() or not peacock_dir.exists():
+        print("Error: Dataset directory must contain 'HEN' and 'PEACOCK' subdirectories.")
+        print(f"Looked in: {args.dataset_dir}")
+        sys.exit(1)
+
+    # Count source images
+    hen_imgs = list(hen_dir.glob("*.[jJ][pP][gG]")) + list(hen_dir.glob("*.[jJ][pP][eE][gG]")) + list(hen_dir.glob("*.[pP][nN][gG]"))
+    peacock_imgs = list(peacock_dir.glob("*.[jJ][pP][gG]")) + list(peacock_dir.glob("*.[jJ][pP][eE][gG]")) + list(peacock_dir.glob("*.[pP][nN][gG]"))
+
+    print(f"Found {len(hen_imgs)} images in HEN/")
+    print(f"Found {len(peacock_imgs)} images in PEACOCK/")
+
+    if len(hen_imgs) == 0 or len(peacock_imgs) == 0:
+        print("Error: Both HEN and PEACOCK folders must contain at least a few images to train.")
+        sys.exit(1)
+
+    # 2. Setup OTHER class automatically
+    setup_other_class(str(other_dir))
+    other_imgs = list(other_dir.glob("*.[jJ][pP][gG]")) + list(other_dir.glob("*.[jJ][pP][eE][gG]")) + list(other_dir.glob("*.[pP][nN][gG]"))
+
+    # 3. Gather paths and labels
+    class_names = ["HEN", "OTHER", "PEACOCK"]  # Alphabetical order
+    class_to_idx = {name: i for i, name in enumerate(class_names)}
+    
+    image_paths = []
+    labels = []
+
+    for name in class_names:
+        dir_path = dataset_path / name
+        imgs = list(dir_path.glob("*.[jJ][pP][gG]")) + list(dir_path.glob("*.[jJ][pP][eE][gG]")) + list(dir_path.glob("*.[pP][nN][gG]")) + list(dir_path.glob("*.[wW][eE][bB][pP]"))
+        for img in imgs:
+            # Simple pre-check to make sure it loads
+            try:
+                with Image.open(img) as temp_img:
+                    temp_img.draft("RGB", (32, 32))
+                image_paths.append(str(img))
+                labels.append(class_to_idx[name])
+            except Exception:
+                print(f"  [Warning] Skipping corrupt file: {img}")
+
+    print(f"Total valid training samples: {len(image_paths)}")
+    for name in class_names:
+        idx = class_to_idx[name]
+        count = labels.count(idx)
+        print(f"  Class '{name}': {count} images")
+
+    # 4. Load online pretrained model and processor
+    print("-> Loading pre-trained base model and image processor...")
+    image_processor = AutoImageProcessor.from_pretrained("chriamue/bird-species-classifier")
+    model = AutoModelForImageClassification.from_pretrained("chriamue/bird-species-classifier")
+
+    # 5. Freeze base layers to make training super fast and prevent catastrophic overfitting
+    print("-> Freezing feature extractor base weights (training classification head only)...")
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # 6. Replace classification head
+    # The original was: (classifier): Linear(in_features=1408, out_features=525, bias=True)
+    num_features = model.classifier.in_features
+    model.classifier = nn.Linear(num_features, len(class_names))
+    
+    # Configure model config metadata so Hugging Face saves labels correctly
+    model.config.id2label = {i: name for i, name in enumerate(class_names)}
+    model.config.label2id = {name: i for i, name in enumerate(class_names)}
+    model.config.num_labels = len(class_names)
+
+    # 7. Create DataLoader
+    dataset = BirdDataset(image_paths, labels, image_processor)
+    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
+
+    # 8. Setup optimizer and loss
+    optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=args.lr)
+    criterion = nn.CrossEntropyLoss()
+
+    # 9. Training Loop
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.train()
+
+    print(f"-> Starting training on device: {device}...")
+    for epoch in range(args.epochs):
+        running_loss = 0.0
+        correct = 0
+        total = 0
+        
+        for batch_x, batch_y in dataloader:
+            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
+            
+            optimizer.zero_grad()
+            outputs = model(pixel_values=batch_x)
+            logits = outputs.logits
+            
+            loss = criterion(logits, batch_y)
+            loss.backward()
+            optimizer.step()
+            
+            running_loss += loss.item() * batch_x.size(0)
+            _, predicted = torch.max(logits, 1)
+            total += batch_y.size(0)
+            correct += (predicted == batch_y).sum().item()
+
+        epoch_loss = running_loss / total
+        epoch_acc = correct / total
+        print(f"  Epoch {epoch+1}/{args.epochs} - Loss: {epoch_loss:.4f} - Accuracy: {epoch_acc:.4f}")
+
+    # 10. Save fine-tuned model
+    print(f"-> Saving fine-tuned model to {args.output_dir}...")
+    os.makedirs(args.output_dir, exist_ok=True)
+    model.save_pretrained(args.output_dir)
+    image_processor.save_pretrained(args.output_dir)
+    print("-> Done! Custom model training completed successfully.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Fine-tune bird species classifier on custom dataset.")
+    parser.add_argument("--dataset_dir", type=str, default="dataset", help="Directory containing HEN and PEACOCK folders")
+    parser.add_argument("--output_dir", type=str, default="custom_model", help="Directory to save custom model weights")
+    parser.add_argument("--epochs", type=int, default=8, help="Number of training epochs")
+    parser.add_argument("--batch_size", type=int, default=8, help="DataLoader batch size")
+    parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate for the classification head")
+    
+    args = parser.parse_args()
+    train_model(args)

verify_custom_model.py

@@ -0,0 +1,85 @@
+import os
+import sys
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+def verify():
+    custom_dir = "custom_model"
+    if not os.path.exists(custom_dir):
+        print(f"Error: custom_model directory '{custom_dir}' does not exist.")
+        sys.exit(1)
+
+    print("-> Loading custom fine-tuned model...")
+    try:
+        processor = AutoImageProcessor.from_pretrained(custom_dir)
+        model = AutoModelForImageClassification.from_pretrained(custom_dir)
+        print("[OK] Custom model loaded successfully!")
+    except Exception as e:
+        print(f"[ERROR] Failed to load custom model: {e}")
+        sys.exit(1)
+
+    # Print labels
+    labels = list(model.config.id2label.values())
+    print(f"Model Labels: {labels}")
+    assert "HEN" in labels, "HEN label missing"
+    assert "PEACOCK" in labels, "PEACOCK label missing"
+    assert "OTHER" in labels, "OTHER label missing"
+
+    # Define test images
+    test_cases = [
+        ("dataset/HEN/hen1.jpg", "HEN"),
+        ("dataset/PEACOCK/peacock1.jpg", "PEACOCK"),
+    ]
+
+    # Find a file in OTHER to test
+    other_dir = "dataset/OTHER"
+    if os.path.exists(other_dir):
+        other_files = [f for f in os.listdir(other_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+        if other_files:
+            test_cases.append((os.path.join(other_dir, other_files[0]), "OTHER"))
+
+    print("\n-> Running prediction tests...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+
+    all_passed = True
+    for img_path, expected_class in test_cases:
+        if not os.path.exists(img_path):
+            print(f"[Warning] Test image '{img_path}' not found. Skipping test.")
+            continue
+
+        try:
+            image = Image.open(img_path).convert("RGB")
+            inputs = processor(images=image, return_tensors="pt").to(device)
+            
+            with torch.no_grad():
+                outputs = model(**inputs)
+            
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=-1)
+            pred_idx = torch.argmax(probs, dim=-1).item()
+            pred_label = model.config.id2label[pred_idx]
+            confidence = probs[0][pred_idx].item()
+
+            print(f"Image: {img_path}")
+            print(f"  Expected: {expected_class}")
+            print(f"  Predicted: {pred_label} (confidence: {confidence:.4f})")
+            
+            if pred_label == expected_class:
+                print("  [PASS]")
+            else:
+                print("  [FAIL] (Mismatch)")
+                all_passed = False
+        except Exception as e:
+            print(f"  [FAIL] (Error: {e})")
+            all_passed = False
+
+    if all_passed:
+        print("\nAll verification tests passed successfully!")
+    else:
+        print("\nSome verification tests failed. Please inspect the outputs.")
+
+if __name__ == "__main__":
+    verify()