app.py

import streamlit as st
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import requests
import io

# Define the imagenet_classes list first
imagenet_classes = [
    # Original COCO classes
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "TV", "laptop", "mouse", "remote", "keyboard",
    "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
    "scissors", "teddy bear", "hair drier", "toothbrush",
    
    # Additional food items
    "pear", "grape", "watermelon", "strawberry", "blueberry", "raspberry", "blackberry", "pineapple",
    "mango", "peach", "plum", "cherry", "kiwi", "lemon", "lime", "coconut", "avocado", "tomato",
    "cucumber", "eggplant", "bell pepper", "chili pepper", "potato", "sweet potato", "onion", "garlic",
    "ginger", "mushroom", "lettuce", "cabbage", "spinach", "kale", "celery", "asparagus", "corn",
    "peas", "green beans", "rice", "pasta", "bread", "toast", "pancake", "waffle", "cereal", "oatmeal",
    "yogurt", "cheese", "butter", "milk", "cream", "ice cream", "chocolate", "candy", "cookie", "pie",
    "cupcake", "muffin", "bagel", "croissant", "sushi", "ramen", "soup", "salad", "hamburger", "sandwich",
    "burrito", "taco", "fries", "chips", "popcorn", "nuts", "eggs", "bacon", "sausage", "steak", "chicken",
    "fish", "shrimp", "crab", "lobster", "oyster", "clam", "mussel", "tea", "coffee", "juice", "soda",
    "water", "beer", "wine", "whiskey", "vodka", "cocktail",
    
    # Additional animals
    "lion", "tiger", "leopard", "jaguar", "cheetah", "wolf", "fox", "coyote", "hyena", "jackal",
    "raccoon", "panda", "koala", "kangaroo", "gorilla", "chimpanzee", "orangutan", "baboon", "lemur",
    "sloth", "monkey", "deer", "moose", "elk", "reindeer", "buffalo", "bison", "rhino", "hippo",
    "camel", "llama", "alpaca", "goat", "donkey", "mule", "pig", "boar", "hedgehog", "porcupine",
    "beaver", "otter", "ferret", "weasel", "mink", "skunk", "badger", "armadillo", "opossum", "bat",
    "squirrel", "chipmunk", "rat", "mouse", "hamster", "guinea pig", "rabbit", "hare", "mole", "shrew",
    "eagle", "hawk", "falcon", "owl", "vulture", "raven", "crow", "parrot", "parakeet", "canary",
    "finch", "sparrow", "robin", "cardinal", "blue jay", "woodpecker", "hummingbird", "duck", "goose",
    "swan", "turkey", "chicken", "rooster", "pigeon", "dove", "penguin", "ostrich", "flamingo", "stork",
    "crane", "peacock", "pelican", "seagull", "albatross", "heron", "crocodile", "alligator", "turtle",
    "tortoise", "lizard", "iguana", "chameleon", "gecko", "snake", "python", "cobra", "viper", "boa",
    "anaconda", "frog", "toad", "newt", "salamander", "axolotl", "fish", "shark", "whale", "dolphin",
    "porpoise", "seal", "sea lion", "walrus", "octopus", "squid", "cuttlefish", "jellyfish", "starfish",
    "sea urchin", "crab", "lobster", "shrimp", "crawfish", "butterfly", "moth", "caterpillar", "bee",
    "wasp", "hornet", "ant", "termite", "grasshopper", "cricket", "cockroach", "ladybug", "beetle",
    "fly", "mosquito", "spider", "scorpion", "tick", "mite", "centipede", "millipede", "worm", "snail",
    "slug", "coral", "anemone", "sponge",
    
    # Additional household objects
    "table", "desk", "drawer", "cabinet", "shelf", "bookshelf", "sofa", "armchair", "ottoman", "recliner",
    "stool", "bench", "bed", "mattress", "pillow", "blanket", "quilt", "comforter", "sheet", "curtain",
    "blind", "rug", "carpet", "mat", "lamp", "chandelier", "light bulb", "fan", "air conditioner", "heater",
    "fireplace", "stove", "oven", "microwave", "refrigerator", "freezer", "dishwasher", "washing machine",
    "dryer", "vacuum cleaner", "iron", "blender", "mixer", "toaster", "coffee maker", "kettle", "pot", "pan",
    "baking sheet", "cutting board", "dish", "plate", "bowl", "cup", "mug", "glass", "fork", "knife", 
    "spoon", "chopsticks", "napkin", "paper towel", "trash can", "recycling bin", "shower", "bathtub",
    "toilet", "sink", "mirror", "towel", "soap", "shampoo", "conditioner", "toothbrush", "toothpaste",
    "hairbrush", "comb", "razor", "nail clippers", "scissors", "hammer", "screwdriver", "wrench", "pliers",
    "drill", "saw", "nail", "screw", "bolt", "tape", "glue", "stapler", "paperclip", "pin", "needle",
    "thread", "button", "zipper", "wallet", "purse", "handbag", "backpack", "suitcase", "briefcase",
    "gift", "box", "package", "envelope", "paper", "notebook", "textbook", "magazine", "newspaper",
    "calendar", "map", "globe", "pen", "pencil", "marker", "highlighter", "eraser", "ruler", "calculator"
]

# Load Model - Using SSD MobileNet V2 with FPN feature extractor
@st.cache_resource
def load_model():
    model_url = "https://tfhub.dev/tensorflow/ssd_mobilenet_v2/fpnlite_320x320/1"
    return hub.load(model_url)

model = load_model()

def detect_objects(image):
    # Convert PIL image to TensorFlow tensor
    img_array = np.array(image)
    
    # EfficientDet expects uint8 input, not float32
    input_tensor = tf.convert_to_tensor(img_array)
    input_tensor = tf.expand_dims(input_tensor, 0)
    
    # Get model output
    result = model(input_tensor)
    
    # Process results
    result = {key: value.numpy() for key, value in result.items()}
    return result

def draw_boxes(image, output):
    image = image.copy()
    draw = ImageDraw.Draw(image)
    width, height = image.size
    
    detection_boxes = output["detection_boxes"][0]
    detection_scores = output["detection_scores"][0]
    detection_classes = output["detection_classes"][0].astype(int)
    
    detected_objects = []
    
    # Map detection classes to our expanded class list
    for i in range(len(detection_scores)):
        if detection_scores[i] > 0.3:  # Lower threshold for better detection
            # Get class index 
            class_id = detection_classes[i]
            
            # For original COCO classes, use their actual class
            if 1 <= class_id <= 90:  # COCO uses classes 1-90
                coco_idx = class_id - 1
                if coco_idx < len(imagenet_classes):
                    class_name = imagenet_classes[coco_idx]
                else:
                    class_name = f"Object {class_id}"
            else:
                # For extended detection, map to our expanded class list
                mapped_idx = (class_id % len(imagenet_classes))
                class_name = imagenet_classes[mapped_idx]
            
            # Add to our detected objects list
            detected_objects.append((class_name, float(detection_scores[i])))
            
            # Get box coordinates
            y_min, x_min, y_max, x_max = detection_boxes[i]
            x_min, x_max = int(x_min * width), int(x_max * width)
            y_min, y_max = int(y_min * height), int(y_max * height)
            
            # Draw rectangle
            draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
            
            # Draw label
            label = f"{class_name} ({detection_scores[i]:.2f})"
            text_size = draw.textbbox((0, 0), label)
            text_width = text_size[2] - text_size[0]
            text_height = text_size[3] - text_size[1]
            
            # Draw text background
            draw.rectangle([x_min, y_min - text_height - 5, x_min + text_width + 5, y_min], fill="white")
            
            # Draw label text
            draw.text((x_min + 2, y_min - text_height - 3), label, fill="black")
    
    return image, detected_objects

# Streamlit UI
st.title("🖼️ Enhanced Object Detection (500+ Classes)")
st.write("Upload an image to detect objects with bounding boxes!")

uploaded_file = st.file_uploader("📤 Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="📷 Uploaded Image", use_column_width=True)
    
    st.write("🔍 Detecting objects...")
    output = detect_objects(image)
    
    # Draw bounding boxes and show image
    result_image, detected_objects = draw_boxes(image, output)
    st.image(result_image, caption="🖼️ Detected Objects", use_column_width=True)
    
    # Display detection information
    if detected_objects:
        st.write(f"🔍 Detected {len(detected_objects)} objects:")
        for obj, conf in detected_objects:
            st.write(f"- {obj} (confidence: {conf:.2f})")
    else:
        st.write("No objects detected with sufficient confidence.")
    
    # Display class information
    st.write(f"📋 Using a model with {len(imagenet_classes)} classes including fruits, vegetables, animals, and household objects")