app.py

import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont

from transformers import AutoImageProcessor
from transformers import AutoModelForObjectDetection

# Note: Can load from Hugging Face or can load from local.
# You will have to replace {mrdbourke} for your own username if the model is on your Hugging Face account.
model_save_path = "mrdbourke/detr_finetuned_trashify_box_detector_with_data_aug" 

# Load the model and preprocessor
image_processor = AutoImageProcessor.from_pretrained(model_save_path)
model = AutoModelForObjectDetection.from_pretrained(model_save_path)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Get the id2label dictionary from the model
id2label = model.config.id2label

# Set up a colour dictionary for plotting boxes with different colours
color_dict = {   
    "bin": "green",
    "trash": "blue",
    "hand": "purple",
    "trash_arm": "yellow",
    "not_trash": "red",
    "not_bin": "red",
    "not_hand": "red",
}

# Create helper functions for seeing if items from one list are in another 
def any_in_list(list_a, list_b):
    "Returns True if any item from list_a is in list_b, otherwise False."
    return any(item in list_b for item in list_a)

def all_in_list(list_a, list_b):
    "Returns True if all items from list_a are in list_b, otherwise False."
    return all(item in list_b for item in list_a)

def filter_highest_scoring_box_per_class(boxes, labels, scores):
    """
    Perform NMS (Non-max Supression) to only keep the top scoring box per class.

    Args:
        boxes: tensor of shape (N, 4)
        labels: tensor of shape (N,)
        scores: tensor of shape (N,)
    Returns:
        boxes: tensor of shape (N, 4) filtered for max scoring item per class
        labels: tensor of shape (N,) filtered for max scoring item per class
        scores: tensor of shape (N,) filtered for max scoring item per class
    """
    # Start with a blank keep mask (e.g. all False and then update the boxes to keep with True)
    keep_mask = torch.zeros(len(boxes), dtype=torch.bool)

    # For each unique class
    for class_id in labels.unique():
        # Get the indicies for the target class
        class_mask = labels == class_id

        # If any of the labels match the current class_id
        if class_mask.any():
            # Find the index of highest scoring box for this specific class
            class_scores = scores[class_mask]
            highest_score_idx = class_scores.argmax()

            # Convert back to the original index
            original_idx = torch.where(class_mask)[0][highest_score_idx]

            # Update the index in the keep mask to keep the highest scoring box 
            keep_mask[original_idx] = True
        
    return boxes[keep_mask], labels[keep_mask], scores[keep_mask]

def create_return_string(list_of_predicted_labels, target_items=["trash", "bin", "hand"]):
     # Setup blank string to print out
    return_string = ""

    # If no items detected or trash, bin, hand not in list, return notification 
    if (len(list_of_predicted_labels) == 0) or not (any_in_list(list_a=target_items, list_b=list_of_predicted_labels)):
        return_string = f"No trash, bin or hand detected at confidence threshold {conf_threshold}. Try another image or lowering the confidence threshold."
        return return_string

    # If there are some missing, print the ones which are missing
    elif not all_in_list(list_a=target_items, list_b=list_of_predicted_labels):
        missing_items = []
        for item in target_items:
            if item not in list_of_predicted_labels:
                missing_items.append(item)
        return_string = f"Detected the following items: {list_of_predicted_labels} (total: {len(list_of_predicted_labels)}). But missing the following in order to get +1: {missing_items}. If this is an error, try another image or altering the confidence threshold. Otherwise, the model may need to be updated with better data."
        
    # If all 3 trash, bin, hand occur = + 1
    if all_in_list(list_a=target_items, list_b=list_of_predicted_labels):
        return_string = f"+1! Found the following items: {list_of_predicted_labels} (total: {len(list_of_predicted_labels)}), thank you for cleaning up the area!"

    print(return_string)

    return return_string

def predict_on_image(image, conf_threshold):
    with torch.no_grad():
        inputs = image_processor(images=[image], return_tensors="pt")
        outputs = model(**inputs.to(device))

        target_sizes = torch.tensor([[image.size[1], image.size[0]]]) # height, width 

        results = image_processor.post_process_object_detection(outputs,
                                                                threshold=conf_threshold,
                                                                target_sizes=target_sizes)[0]
    # Return all items in results to CPU
    for key, value in results.items():
        try:
            results[key] = value.item().cpu() # can't get scalar as .item() so add try/except block
        except:
            results[key] = value.cpu()

    # Can return results as plotted on a PIL image (then display the image)
    draw = ImageDraw.Draw(image)

    # Create a copy of the image to draw on it for NMS
    image_nms = image.copy()
    draw_nms = ImageDraw.Draw(image_nms)

    # Get a font from ImageFont
    font = ImageFont.load_default(size=20)

    # Get class names as text for print out
    class_name_text_labels = []

    # TK - update this for NMS
    class_name_text_labels_nms = []

    # Get original boxes, scores, labels
    original_boxes = results["boxes"]
    original_labels = results["labels"]
    original_scores = results["scores"]

    # Filter boxes and only keep 1x of each label with highest score
    filtered_boxes, filtered_labels, filtered_scores = filter_highest_scoring_box_per_class(boxes=original_boxes,
                                                                                            labels=original_labels,
                                                                                            scores=original_scores)
    # TODO: turn this into a function so it's cleaner?
    for box, label, score in zip(original_boxes, original_labels, original_scores):
        # Create coordinates
        x, y, x2, y2 = tuple(box.tolist())

        # Get label_name
        label_name = id2label[label.item()]
        targ_color = color_dict[label_name]
        class_name_text_labels.append(label_name)

        # Draw the rectangle
        draw.rectangle(xy=(x, y, x2, y2), 
                       outline=targ_color,
                       width=3)
        
        # Create a text string to display
        text_string_to_show = f"{label_name} ({round(score.item(), 3)})"

        # Draw the text on the image
        draw.text(xy=(x, y),
                  text=text_string_to_show,
                  fill="white",
                  font=font)
    
    # TODO: turn this into a function so it's cleaner?
    for box, label, score in zip(filtered_boxes, filtered_labels, filtered_scores):
        # Create coordinates
        x, y, x2, y2 = tuple(box.tolist())

        # Get label_name
        label_name = id2label[label.item()]
        targ_color = color_dict[label_name]
        class_name_text_labels_nms.append(label_name)

        # Draw the rectangle
        draw_nms.rectangle(xy=(x, y, x2, y2), 
                       outline=targ_color,
                       width=3)
        
        # Create a text string to display
        text_string_to_show = f"{label_name} ({round(score.item(), 3)})"

        # Draw the text on the image
        draw_nms.text(xy=(x, y),
                  text=text_string_to_show,
                  fill="white",
                  font=font)
    
    
    # Remove the draw each time
    del draw
    del draw_nms

    # Create the return string
    return_string = create_return_string(list_of_predicted_labels=class_name_text_labels)
    return_string_nms = create_return_string(list_of_predicted_labels=class_name_text_labels_nms)
    
    return image, return_string, image_nms, return_string_nms

# Create the interface
demo = gr.Interface(
    fn=predict_on_image,
    inputs=[
        gr.Image(type="pil", label="Target Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold")
    ],
    outputs=[
        gr.Image(type="pil", label="Image Output (no filtering)"),
        gr.Text(label="Text Output (no filtering)"),
        gr.Image(type="pil", label="Image Output (with max score per class box filtering)"),
        gr.Text(label="Text Output (with max score per class box filtering)")
        
    ],
    title="🚮 Trashify Object Detection Demo V3",
    description="""Help clean up your local area! Upload an image and get +1 if there is all of the following items detected: trash, bin, hand.

    The model in V3 is [same model](https://huggingface.co/mrdbourke/detr_finetuned_trashify_box_detector_with_data_aug) as in [V2](https://huggingface.co/spaces/mrdbourke/trashify_demo_v2) (trained with data augmentation) but has an additional post-processing step (NMS or [Non Maximum Suppression](https://paperswithcode.com/method/non-maximum-suppression)) to filter classes for only the highest scoring box of each class. 
    """,
    # Examples come in the form of a list of lists, where each inner list contains elements to prefill the `inputs` parameter with
    examples=[
        ["examples/trashify_example_1.jpeg", 0.25],
        ["examples/trashify_example_2.jpeg", 0.25],
        ["examples/trashify_example_3.jpeg", 0.25]
    ],
    cache_examples=True
)

# Launch the demo
demo.launch()