Update app.py

app.py

@@ -2,9 +2,8 @@ import random
 import os
 import numpy as np
 import gradio as gr
-from huggingface_hub import from_pretrained_fastai
 from PIL import Image
-from groundingdino.util.inference import load_model
+from groundingdino.util.inference import load_model as load_groundingdino_model
 from groundingdino.util.inference import predict as grounding_dino_predict
 import groundingdino.datasets.transforms as T
 import torch
@@ -13,9 +12,124 @@ from torchvision.transforms.functional import to_tensor
 from torchvision.transforms import GaussianBlur
 import time
 
-from Ambrosia import pre_process_image
+# ----------------------------
+# DINOv2 Classifier Imports
+# ----------------------------
+import torch.nn as nn
+from torchvision import transforms
+import pandas as pd
+from typing import List, Tuple
+import copy
+import matplotlib.pyplot as plt
 
+# ----------------------------
+# DINOv2 Classifier Definitions
+# ----------------------------
 
+# 1. PadToSquare Class
+class PadToSquare:
+    """
+    Pads an image to make it square by adding padding to the shorter side.
+    """
+    def __init__(self, fill=0):
+        self.fill = fill
+
+    def __call__(self, img):
+        w, h = img.size
+        max_wh = max(w, h)
+        hp = (max_wh - w) // 2
+        vp = (max_wh - h) // 2
+        padding = (hp, vp, max_wh - w - hp, max_wh - h - vp)
+        return transforms.functional.pad(img, padding, fill=self.fill, padding_mode='constant')
+
+# 2. DinoVisionTransformerClassifier Class (Modified to include entropy-based approach)
+class DinoVisionTransformerClassifier(nn.Module):
+    """
+    DINOv2 Vision Transformer-based classifier with entropy-based "Unknown" class handling.
+    """
+    def __init__(self, num_classes, hidden_size=256, dropout_p=0.5, negative_slope=0.01):
+        super(DinoVisionTransformerClassifier, self).__init__()
+        # Load DINOv2 model from torch.hub
+        self.transformer = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14', pretrained=True)
+        self.transformer.norm = nn.Identity()  # Remove existing normalization if necessary
+
+        # Batch Normalization after transformer
+        self.batch_norm1 = nn.BatchNorm1d(384)  # 384 is the embedding size
+
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.Linear(384, hidden_size),
+            nn.BatchNorm1d(hidden_size),
+            nn.LeakyReLU(negative_slope=negative_slope, inplace=True),
+            nn.Dropout(p=dropout_p),
+            nn.Linear(hidden_size, num_classes)
+        )
+
+        # Initialize weights
+        self._initialize_weights()
+
+    def forward(self, x):
+        features = self.transformer(x)           # Forward pass through the transformer
+        features = self.batch_norm1(features)    # Apply Batch Normalization
+        logits = self.classifier(features)       # Forward pass through the classification head
+        return logits, features  # Return both logits and features
+
+    def _initialize_weights(self):
+        # Initialize weights of the classifier layers
+        for m in self.classifier.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, a=0.01, mode='fan_in', nonlinearity='leaky_relu')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+
+# 3. Model Loading Function (Updated for Entropy-Based Classifier)
+def load_model(model_path, device):
+    """
+    Loads the trained model and class information from the saved checkpoint.
+
+    Args:
+        model_path (str): Path to the saved .pth model file.
+        device (torch.device): Device to load the model onto.
+
+    Returns:
+        model (nn.Module): The loaded PyTorch model.
+        class_names (List[str]): List of class names.
+    """
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file '{model_path}' does not exist.")
+
+    checkpoint = torch.load(model_path, map_location=device)
+    class_names = checkpoint['class_names']
+    num_classes = len(class_names)
+
+    # Initialize the model architecture
+    model = DinoVisionTransformerClassifier(num_classes=num_classes)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.to(device)
+    model.eval()  # Set to evaluation mode
+
+    return model, class_names
+
+# 4. Image Preprocessing Function (Updated to accept PIL Image directly)
+def preprocess_image_pil(pil_image: Image.Image, transform: transforms.Compose) -> torch.Tensor:
+    """
+    Applies the transformation pipeline to a PIL image.
+
+    Args:
+        pil_image (PIL.Image.Image): The image to preprocess.
+        transform (transforms.Compose): The transformation pipeline.
+
+    Returns:
+        torch.Tensor: The preprocessed image tensor.
+    """
+    return transform(pil_image)
+
+# ----------------------------
+# Gradio App Definitions
+# ----------------------------
 
 # Automatically set device based on availability
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -23,8 +137,7 @@ print(f"Using device: {DEVICE}")
 
 PROMPT = "bug"
 
-
-# Define a custom transform for Gaussian blur
+# Define a custom transform for Gaussian blur (Unused in current context)
 def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1, sigma_max=3):
     if x.ndim == 4:
         for i in range(x.shape[0]):
@@ -34,39 +147,13 @@ def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1
                 x[i] = GaussianBlur(kernel_size=kernel_size, sigma=sigma)(x[i])
     return x
 
-# Custom Label Function
+# Custom Label Function (Unused in current context)
 def custom_label_func(fpath):
     # this directs the labels to be 2 levels up from the image folder
     label = fpath.parents[2].name
     return label
 
-# this function only describes how much a singular value in al ist stands out.
-# if all values in the lsit are high or low this is 1
-# the smaller the proportiopn of number of disimilar vlaues are to other more similar values the lower this number
-# the larger the gap between the dissimilar numbers and the simialr number the smaller this number
-# only able to interpret probabilities or values between 0 and 1
-# this function outputs an estimate an inverse of the classification confidence based on the probabilities of all the classes.
-# the wedge threshold splits the data on a threshold with a magnitude of a positive int to force a ledge/peak in the data
-def unkown_prob_calc(probs, wedge_threshold, wedge_magnitude=1, wedge='strict'):
-    if wedge =='strict':
-        increase_var = (1/(wedge_magnitude))
-        decrease_var = (wedge_magnitude)
-    if wedge =='dynamic': # this allows pointsthat are furhter from the threshold ot be moved less and points clsoer to be moved more
-        increase_var = (1/(wedge_magnitude*((1-np.abs(probs-wedge_threshold)))))
-        decrease_var = (wedge_magnitude*((1-np.abs(probs-wedge_threshold))))
-    else:
-        print("Error: use 'strict' (default) or 'dynamic' as options for the wedge parameter!")
-    probs = np.where(probs>=wedge_threshold , probs**increase_var, probs)
-    probs = np.where(probs<=wedge_threshold , probs**decrease_var, probs)
-    diff_matrix = np.abs(probs[:, np.newaxis] - probs)
-    diff_matrix_sum = np.sum(diff_matrix)
-    probs_sum = np.sum(probs)
-    class_val = (diff_matrix_sum/probs_sum)
-    max_class_val = ((len(probs)-1)*2)
-    kown_prob = class_val/max_class_val
-    unknown_prob = 1-kown_prob
-    return(unknown_prob)
-
+# Image loading function for GroundingDINO
 def load_image(image_source):
     transform = T.Compose(
         [
@@ -80,17 +167,47 @@ def load_image(image_source):
     image_transformed, _ = transform(image_source, None)
     return image_transformed
 
-# load object detection model
-od_model = load_model(
+# Load GroundingDINO object detection model
+od_model = load_groundingdino_model(
     model_checkpoint_path="groundingdino_swint_ogc.pth", 
     model_config_path="GroundingDINO_SwinT_OGC.cfg.py", 
     device=DEVICE)
 print("Object detection model loaded")
 
+# Load DINOv2 classifier model (Updated to use the entropy-based classifier)
+# Update MODEL_PATH to the path where your DINOv2 model checkpoint is stored
+MODEL_PATH = 'dinov2_classifier_with_vos_unsure.pth'  # Updated model path
+dinov2_model, class_names = load_model(MODEL_PATH, torch.device(DEVICE))
+print(f"DINOv2 Classification model loaded with {len(class_names)} classes.")
+
+# Optionally, append "Unknown" to class names if needed
+# Removed the line that appends "Unknown" as the model handles it via thresholding
+
+# Replace specific class names if necessary
+# Example: Replace "Scolotodes_schwarzi" with "Scolytodes_glaber"
+target = "Scolotodes_schwarzi"
+if target in class_names:
+    idx = class_names.index(target)
+    class_names[idx] = "Scolytodes_glaber"
+    print(f"Replaced '{target}' with 'Scolytodes_glaber' in class names.")
+else:
+    print(f"'{target}' not found in class names. No replacement made.")
+
+# Define the transformation pipeline for DINOv2 model
+dinov2_transform = transforms.Compose([
+    transforms.Resize(224),            # Resize smaller edge to 224
+    PadToSquare(),                     # Pad to make the image square
+    transforms.Resize((224, 224)),      # Resize to 224x224
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],  # Normalize with ImageNet mean
+                         [0.229, 0.224, 0.225])  # Normalize with ImageNet std
+])
+
+# Object Detection Function
 def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"):
     TEXT_PROMPT = prompt
-    BOX_TRESHOLD = 0.35
-    TEXT_TRESHOLD = 0.25
+    BOX_THRESHOLD = 0.35  # Adjusted back to original value
+    TEXT_THRESHOLD = 0.25  # Adjusted back to original value
     DEVICE = device  # cuda or cpu
     
     # Convert numpy array to PIL Image if needed
@@ -102,18 +219,18 @@ def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"
     # Transform the image
     image_transformed = load_image(image_source = og_image_obj)
     
-    # Your model prediction code here...
+    # Model prediction
     boxes, logits, phrases = grounding_dino_predict(
         model=model,
         image=image_transformed,
         caption=TEXT_PROMPT,
-        box_threshold=BOX_TRESHOLD,
-        text_threshold=TEXT_TRESHOLD,
+        box_threshold=BOX_THRESHOLD,
+        text_threshold=TEXT_THRESHOLD,
         device=DEVICE)
 
     # Use og_image_obj directly for further processing
-    height, width = og_image_obj.size
-    boxes_norm = boxes * torch.Tensor([height, width, height, width])
+    width, height = og_image_obj.size  # Corrected to (width, height)
+    boxes_norm = boxes * torch.Tensor([width, height, width, height])
     xyxy = box_convert(
         boxes=boxes_norm,
         in_fmt="cxcywh",
@@ -122,111 +239,139 @@ def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"
     for i in range(len(boxes_norm)):
         crop_img = og_image_obj.crop((xyxy[i]))
         img_lst.append(crop_img)
-    return (img_lst)
-
-
-# load beetle classifier model
-repo_id="ChristopherMarais/beetle-model-mini"
-bc_model = from_pretrained_fastai(repo_id)
-bc_model.to(DEVICE)
-# get class names
-labels = np.append(np.array(bc_model.dls.vocab), "Unknown")
-# Replace some names used in the classifier
-# Check if the element was found to prevent errors
-# The target value you're looking for
-target = "Scolotodes_schwarzi"
-# Finding the index using np.where
-indices = np.where(labels == target)
-# Extracting the first occurrence, if found
-if indices[0].size > 0:
-    idx = indices[0][0]
-    print(f"Index of {target}: {idx}")
-else:
-    print(f"{target} not found in the array.")
-# Replace occurence
-if idx != -1:
-    labels[idx] = "Scolytodes_glaber"
-print("Classification model loaded")
+    print(f"Detected {len(img_lst)} objects.")
+    return img_lst
+
+# Inference/Class Prediction Function using the Entropy-Based DINOv2 Classifier
+def classify_beetle(img: Image.Image, threshold=75.0):
+    """
+    Classifies the input image using the DINOv2 classifier with entropy-based "Unknown" class.
+
+    Args:
+        img (PIL.Image.Image): The image to classify.
+        threshold (float): Confidence threshold to assign "Unknown".
+
+    Returns:
+        dict: The top 3 class labels with their corresponding confidence scores and "Unknown" if applicable.
+    """
+    # Preprocess the image
+    input_tensor = preprocess_image_pil(img, dinov2_transform).unsqueeze(0).to(torch.device(DEVICE))
+    print(f"Input tensor shape: {input_tensor.shape}")
+    
+    with torch.no_grad():
+        outputs, _ = dinov2_model(input_tensor)
+        print(f"Model outputs: {outputs}")
+        probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]  # p(x) in [0,1]
+        print(f"Probabilities (0-1 scale): {probabilities}")
+    
+    # Calculate entropy
+    # Adding a small epsilon to avoid log(0)
+    epsilon = 1e-12
+    entropy = -np.sum(probabilities * np.log(probabilities + epsilon))
+    # Maximum entropy for uniform distribution
+    max_entropy = -np.sum((1.0 / len(probabilities)) * np.log(1.0 / len(probabilities)))
+    normalized_entropy = entropy / max_entropy  # Normalize between 0 and 1
+    unknown_prob = normalized_entropy
+    print(f"Entropy: {entropy}, Normalized Entropy: {normalized_entropy}, Unknown Probability: {unknown_prob}")
+    
+    # Convert probabilities to percentage for display
+    probabilities_percent = np.around(probabilities * 100, decimals=1)
+    print(f"Probabilities (Percentage): {probabilities_percent}")
+    
+    # Get top 3 classes
+    top_indices = np.argsort(probabilities_percent)[-3:][::-1]  # Indices of top 3 classes
+    top_probs = probabilities_percent[top_indices]
+    top_classes = [class_names[i] for i in top_indices]
+    
+    # Initialize conf_dict with top 3 classes
+    conf_dict = {top_classes[i]: float(top_probs[i]) for i in range(len(top_classes))}
+    
+    # Assign "Unknown" based on entropy and threshold
+    if top_probs[0] < threshold:
+        conf_dict["Unknown"] = float(np.around(unknown_prob, decimals=1))
+    
+    print(f"Conf_dict: {conf_dict}")
+    
+    return conf_dict
 
+# Main Prediction Function for Gradio
 def predict_beetle(img):
-    print("Detecting & classifying beetles...")
-    start_time = time.perf_counter() # record how long it processes
-    # Split image into smaller images of detected objects
+    print("Detecting objects in the image...")
+    start_time = time.perf_counter()  # Start timing
+
+    # Detect objects in the image
     image_lst = detect_objects(og_image=img, model=od_model, prompt=PROMPT, device=DEVICE)
     
-    # pre_process = pre_process_image(manual_thresh_buffer=0.15, image = img) # use image_dir if directory of image used
-    # pre_process.segment(cluster_num=2, 
-    #                     image_edge_buffer=50)
-    # image_lst = pre_process.col_image_lst
+    print(f"Detected {len(image_lst)} objects.")
     
-    print("Objects detected")
-    end_time = time.perf_counter()
-    processing_time = end_time - start_time
-    print(f"Processing duration: {processing_time} seconds")
-    # get predictions for all segments
-    conf_dict_lst = []
+    # Initialize lists to hold results
     output_lst = []
     img_cnt = len(image_lst)
-    for i in range(0,img_cnt):
-        prob_ar = np.array(bc_model.predict(image_lst[i])[2].to(DEVICE).cpu())
-        unkown_prob = unkown_prob_calc(probs=prob_ar, wedge_threshold=0.85, wedge_magnitude=5, wedge='dynamic')
-        prob_ar = np.append(prob_ar, unkown_prob)
-        prob_ar = np.around(prob_ar*100, decimals=1)
-        # only show the top 5 predictions
-        # Sorting the dictionary by value in descending order and taking the top items
-        top_num = 3
-        conf_dict = {labels[i]: float(prob_ar[i]) for i in range(len(prob_ar))}
-        print(conf_dict)
-        conf_dict = dict(sorted(conf_dict.items(), key=lambda item: item[1], reverse=True)[:top_num])
-        conf_dict_lst.append(str(conf_dict)[1:-1]) # remove dictionary brackets
-        result = list(zip(image_lst, conf_dict_lst))
-        print(f"Beetle classified - {i}")
-        # record how long classification takes
-        end_time = time.perf_counter()
-        processing_time = end_time - start_time
-        print(f"Processing duration: {processing_time} seconds")
-    return(result)
-
-
-# gradio app
-# css = """
-# button {
-#     width: auto;  /* Set your desired width */
-# }
-# """
+    
+    for i in range(img_cnt):
+        print(f"Classifying object {i+1}/{img_cnt}...")
+        conf_dict = classify_beetle(image_lst[i])
+        output_lst.append([image_lst[i], conf_dict])
+        print(f"Object {i+1} classified.")
+    
+    end_time = time.perf_counter()
+    processing_time = end_time - start_time
+    print(f"Total processing duration: {processing_time:.2f} seconds")
+    
+    return output_lst
+
+# ----------------------------
+# Gradio Interface Setup
+# ----------------------------
 
 sample_images_dir = "example_images"
 
 # Sample images with labels
 example_images = [
-    [os.path.join(sample_images_dir, "example1.jpg")],
-    [os.path.join(sample_images_dir, "example2.jpg")],
-    [os.path.join(sample_images_dir, "example3.jpg")],
-    [os.path.join(sample_images_dir, "mixed.jpg")]
-    
+    os.path.join(sample_images_dir, "example1.jpg"),
+    os.path.join(sample_images_dir, "example2.jpg"),
+    os.path.join(sample_images_dir, "example3.jpg"),
+    os.path.join(sample_images_dir, "mixed.jpg")
 ]
 # Corresponding labels for the example images
-example_labels = ["Example Beetle 1", "Example Beetle 2", "Example Beetle 3", "Example Beetles Mixed"]
+example_labels = ["Example Beetles 1", "Example Beetles 2", "Example Beetles 3", "Example Beetles 4"]
 
-with gr.Blocks() as demo: # css=css in the brackets
+with gr.Blocks() as demo:
     gr.Markdown("<h1><center>Bark Beetle Classification</center></h1>")
-    # gr.Markdown("<h3><center>Note this instance of the classifier is for demonstration only and runs on CPU, not on GPU. If you are interested in testing the model, contact us, and we will switch it to its full capacity in an instant.<h3><center>")
+    
     with gr.Column(variant="panel"):
         with gr.Row(variant="compact"):
             inputs = gr.Image(label="Input Image")
-            # Use the `full_width` parameter directly
-            btn = gr.Button("Classify")
+        # Add examples with labels
+        gr.Examples(
+            label="Select an example below if you have no images to upload.", 
+            examples=example_images,
+            inputs=inputs,
+            examples_per_page=4,
+            example_labels=example_labels
+        )
+        
+        btn = gr.Button("Classify", variant="primary")
 
         # Set the gallery layout and height directly in the constructor
-        gallery = gr.Gallery(label="Show images", show_label=True, elem_id="gallery", columns=8, height="auto")
-
-    # Add examples with labels
-    gr.Examples(
-        examples=example_images,
-        inputs=inputs,
-        examples_per_page=4,
-        example_labels=example_labels
-    )
+        gallery = gr.Gallery(label="Classified Objects", show_label=True, elem_id="gallery", columns=4, height="auto")
     
-    btn.click(predict_beetle, inputs, gallery)
-demo.launch(debug=True, show_error=True)
+    # Define the output format for the gallery
+    def format_gallery(results):
+        formatted = []
+        for img, conf in results:
+            # Create a label string from the confidence dictionary
+            label_str = ", ".join([f"{k}: {v:.1f}%" for k, v in conf.items()])
+            # Append the image and label as a tuple
+            formatted.append((img, label_str))
+        return formatted
+
+    # Modify the click event to format the gallery
+    btn.click(
+        lambda img: format_gallery(predict_beetle(img)),
+        inputs,
+        gallery
+    )
+
+# Launch the Gradio app
+demo.launch(share=True, inline=True, debug=True, show_error=True)