Spaces:

souvikg544
/

quickdraw-classifier

Build error

added word

760fccf over 1 year ago

12.7 kB

import gradio as gr
import torch
import clip
from PIL import Image
import numpy as np
import json
from typing import List, Tuple
import io
import base64

# Load CLIP model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

# Load categories
def load_categories():
    """Load drawing categories from categories.txt file"""
    try:
        with open("categories.txt", "r") as f:
            categories = [line.strip() for line in f.readlines()]
        return categories
    except FileNotFoundError:
        # Fallback to a subset of common categories if file not found
        return [
            "airplane", "apple", "car", "cat", "dog", "flower", "house", "tree",
            "sun", "moon", "star", "bird", "fish", "butterfly", "bicycle", "boat"
        ]

# Load categories
CATEGORIES = load_categories()

def preprocess_drawing(image):
    """Preprocess the drawing image for CLIP"""
    if image is None:
        return None
    print(image)
    # Convert to PIL Image if it's not already
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    # print(image)
    # Convert to RGB if it's not already
    if image.mode != 'RGB':
        image = image.convert('RGB')
    # Save the image to a temporary file for debugging
    image.save("temp_drawing.png")
    # Preprocess for CLIP
    return preprocess(image).unsqueeze(0).to(device)

def classify_drawing(image, top_k: int = 5) -> List[Tuple[str, float]]:
    """
    Classify a quickdraw-style drawing using CLIP zero-shot classification
    
    Args:
        image: PIL Image or numpy array of the drawing
        top_k: Number of top predictions to return
    
    Returns:
        List of tuples (category, confidence_score)
    """
    if image is None:
        print("Here")
        return []
    
    try:
        # Preprocess the image
        if isinstance(image, dict):
            print("Image is a dictionary")
            # return []
            print(image)
            image_input = preprocess_drawing(image['composite'])
        else:
            image_input = preprocess_drawing(image)
        if image_input is None:
            # print("Here2")
            return []
        
        # Create text prompts for all categories
        text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
        text_inputs = clip.tokenize(text_prompts).to(device)
        
        # Get features
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            text_features = model.encode_text(text_inputs)
            
            # Normalize features
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            
            # Calculate similarities
            similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
            values, indices = similarities[0].topk(top_k)
        
        # Format results
        results = []
        for i in range(top_k):
            category = CATEGORIES[indices[i]]
            confidence = float(values[i])
            results.append((category, confidence))
        
        return results
    
    except Exception as e:
        print(f"Error in classification: {str(e)}")
        return []
    
def classify_drawing1(image, top_k: int = 5,word: str = "dog") -> List[Tuple[str, float]]:
    """
    Classify a quickdraw-style drawing using CLIP zero-shot classification
    
    Args:
        image: PIL Image or numpy array of the drawing
        top_k: Number of top predictions to return
    
    Returns:
        List of tuples (category, confidence_score)
    """
    if image is None:
        print("Here")
        return []
    
    try:
        # Preprocess the image
        if isinstance(image, dict):
            print("Image is a dictionary")
            # return []
            print(image)
            image_input = preprocess_drawing(image['composite'])
        else:
            image_input = preprocess_drawing(image)
        if image_input is None:
            # print("Here2")
            return []
        if word.lower() not in [cat.lower() for cat in CATEGORIES]:
            CATEGORIES.append(word)
            # print(f"Added new category: {word}")
        
        # Create text prompts for all categories
        text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
        text_inputs = clip.tokenize(text_prompts).to(device)
        
        # Get features
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            text_features = model.encode_text(text_inputs)
            
            # Normalize features
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            
            # Calculate similarities
            similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
            values, indices = similarities[0].topk(top_k)
        
        # Format results
        results = []
        for i in range(top_k):
            category = CATEGORIES[indices[i]]
            confidence = float(values[i])
            results.append((category, confidence))

        if word.lower() not in [cat.lower() for cat, _ in results]:
            # Find the original confidence score for the word
            word_idx = CATEGORIES.index(word)
            word_confidence = float(similarities[0][word_idx])
            
            # Remove the last result if we have top_k results
            if len(results) >= top_k:
                results.pop()
            # Add the word with its original confidence score
            results.append((word, word_confidence))
        
        return results
    
    except Exception as e:
        print(f"Error in classification: {str(e)}")
        return []

def gradio_classify(image):
    """Gradio interface function"""
    if image is None:
        return "Please draw something first!"
    
    results = classify_drawing(image, top_k=5)
    
    if not results:
        return "Sorry, couldn't classify the drawing. Please try again."
    
    # Format results for display
    output = "🎨 **Drawing Classification Results:**\n\n"
    for i, (category, confidence) in enumerate(results, 1):
        percentage = confidence * 100
        output += f"{i}. **{category.title()}** - {percentage:.1f}%\n"
    
    return output

# Create Gradio interface
def create_interface():
    """Create the Gradio interface with both drawing and API tabs"""
    
    with gr.Blocks(
        title="QuickDraw Classifier",
        theme=gr.themes.Soft(),
        css="""
        .main-header { text-align: center; margin-bottom: 2rem; }
        .description { text-align: center; margin-bottom: 1.5rem; font-size: 1.1rem; }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="main-header">
            <h1>🎨 QuickDraw Classifier</h1>
            <p class="description">
                Draw anything and let AI guess what it is! This app uses CLIP for zero-shot classification 
                of your drawings across 345+ categories from Google's QuickDraw dataset.
            </p>
        </div>
        """)
        
        with gr.Tabs():
            # Main Drawing Interface Tab
            with gr.Tab("🖊️ Draw & Classify"):
                with gr.Row():
                    with gr.Column(scale=1):
                        # Drawing canvas
                        canvas = gr.Sketchpad(
                            label="Draw here! 🖊️",
                            type="pil",
                            image_mode="RGB",
                            canvas_size=(400, 400),
                            brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=5),
                            show_download_button=False,
                            show_share_button=False,
                        )
                        
                        with gr.Row():
                            classify_btn = gr.Button("🔍 Classify Drawing", variant="primary", size="lg")
                            clear_btn = gr.ClearButton(canvas, value="🗑️ Clear Canvas", size="lg")
                    
                    with gr.Column(scale=1):
                        # Results
                        output = gr.Markdown(
                            label="Classification Results",
                            value="👈 Draw something and click 'Classify Drawing' to see results!"
                        )
                
                # Examples section
                gr.HTML("""
                <div style="margin-top: 2rem; text-align: center;">
                    <h3>💡 Drawing Tips:</h3>
                    <p>• Draw simple, clear shapes<br>
                    • Try common objects like animals, vehicles, or everyday items<br>
                    • The AI knows 345+ categories including: airplane, cat, house, tree, car, flower, etc.</p>
                </div>
                """)
                
                # Event handlers
                classify_btn.click(
                    fn=gradio_classify,
                    inputs=[canvas],
                    outputs=[output]
                )
                
                # Auto-classify on canvas change (optional, can be removed if too frequent)
                canvas.change(
                    fn=gradio_classify,
                    inputs=[canvas],
                    outputs=[output]
                )
            
            # API Interface Tab
            with gr.Tab("🔌 API Testing"):
                gr.HTML("""
                <div style="text-align: center; margin-bottom: 1rem;">
                    <h3>Base64 Image Classification API</h3>
                    <p>Test the API endpoint with base64 encoded images</p>
                </div>
                """)
                
                with gr.Row():
                    with gr.Column():
                        api_input = gr.Textbox(
                            label="Base64 Image Data", 
                            placeholder="Paste base64 encoded image here...",
                            lines=3
                        )
                        api_top_k = gr.Slider(
                            minimum=1, 
                            maximum=10, 
                            value=5, 
                            step=1, 
                            label="Top K Predictions"
                        )
                        api_word = gr.Textbox(
                            label="Word to Add",
                            placeholder="Enter a word to add to the categories",
                            lines=1
                        )
                        api_btn = gr.Button("🚀 Classify", variant="primary")
                    
                    with gr.Column():
                        api_output = gr.JSON(label="API Response")
                
                api_btn.click(
                    fn=classify_image_api,
                    inputs=[api_input, api_top_k, api_word],
                    outputs=[api_output]
                )
    
    return demo

# API endpoint function for external use
def classify_image_api(image_data: str, top_k: int = 5, word: str = "dog"):
    """
    API endpoint for classifying base64 encoded images
    
    Args:
        image_data: Base64 encoded image string
        top_k: Number of top predictions to return
    
    Returns:
        JSON response with classification results
    """
    try:
        # Decode base64 image
        image_bytes = base64.b64decode(image_data)
        image = Image.open(io.BytesIO(image_bytes))
        
        # Classify
        results = classify_drawing1(image, top_k, word)
        # Normalize confidence scores to sum to 100%
        total_conf = sum(conf for _, conf in results)
        if total_conf > 0:
            results = [(cat, (conf/total_conf) * 100) for cat, conf in results]
        
        return {
            "success": True,
            "predictions": [
                {"category": cat, "confidence": conf} 
                for cat, conf in results
            ]
        }
    
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }

if __name__ == "__main__":
    # Create and launch the main interface
    demo = create_interface()
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    ) 

	import gradio as gr
	import torch
	import clip
	from PIL import Image
	import numpy as np
	import json
	from typing import List, Tuple
	import io
	import base64

	# Load CLIP model
	device = "cuda" if torch.cuda.is_available() else "cpu"
	model, preprocess = clip.load("ViT-B/32", device=device)

	# Load categories
	def load_categories():
	"""Load drawing categories from categories.txt file"""
	try:
	with open("categories.txt", "r") as f:
	categories = [line.strip() for line in f.readlines()]
	return categories
	except FileNotFoundError:
	# Fallback to a subset of common categories if file not found
	return [
	"airplane", "apple", "car", "cat", "dog", "flower", "house", "tree",
	"sun", "moon", "star", "bird", "fish", "butterfly", "bicycle", "boat"
	]

	# Load categories
	CATEGORIES = load_categories()

	def preprocess_drawing(image):
	"""Preprocess the drawing image for CLIP"""
	if image is None:
	return None
	print(image)
	# Convert to PIL Image if it's not already
	if isinstance(image, np.ndarray):
	image = Image.fromarray(image)
	# print(image)
	# Convert to RGB if it's not already
	if image.mode != 'RGB':
	image = image.convert('RGB')
	# Save the image to a temporary file for debugging
	image.save("temp_drawing.png")
	# Preprocess for CLIP
	return preprocess(image).unsqueeze(0).to(device)

	def classify_drawing(image, top_k: int = 5) -> List[Tuple[str, float]]:
	"""
	Classify a quickdraw-style drawing using CLIP zero-shot classification

	Args:
	image: PIL Image or numpy array of the drawing
	top_k: Number of top predictions to return

	Returns:
	List of tuples (category, confidence_score)
	"""
	if image is None:
	print("Here")
	return []

	try:
	# Preprocess the image
	if isinstance(image, dict):
	print("Image is a dictionary")
	# return []
	print(image)
	image_input = preprocess_drawing(image['composite'])
	else:
	image_input = preprocess_drawing(image)
	if image_input is None:
	# print("Here2")
	return []

	# Create text prompts for all categories
	text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
	text_inputs = clip.tokenize(text_prompts).to(device)

	# Get features
	with torch.no_grad():
	image_features = model.encode_image(image_input)
	text_features = model.encode_text(text_inputs)

	# Normalize features
	image_features = image_features / image_features.norm(dim=-1, keepdim=True)
	text_features = text_features / text_features.norm(dim=-1, keepdim=True)

	# Calculate similarities
	similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
	values, indices = similarities[0].topk(top_k)

	# Format results
	results = []
	for i in range(top_k):
	category = CATEGORIES[indices[i]]
	confidence = float(values[i])
	results.append((category, confidence))

	return results

	except Exception as e:
	print(f"Error in classification: {str(e)}")
	return []

	def classify_drawing1(image, top_k: int = 5,word: str = "dog") -> List[Tuple[str, float]]:
	"""
	Classify a quickdraw-style drawing using CLIP zero-shot classification

	Args:
	image: PIL Image or numpy array of the drawing
	top_k: Number of top predictions to return

	Returns:
	List of tuples (category, confidence_score)
	"""
	if image is None:
	print("Here")
	return []

	try:
	# Preprocess the image
	if isinstance(image, dict):
	print("Image is a dictionary")
	# return []
	print(image)
	image_input = preprocess_drawing(image['composite'])
	else:
	image_input = preprocess_drawing(image)
	if image_input is None:
	# print("Here2")
	return []
	if word.lower() not in [cat.lower() for cat in CATEGORIES]:
	CATEGORIES.append(word)
	# print(f"Added new category: {word}")

	# Create text prompts for all categories
	text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
	text_inputs = clip.tokenize(text_prompts).to(device)

	# Get features
	with torch.no_grad():
	image_features = model.encode_image(image_input)
	text_features = model.encode_text(text_inputs)

	# Normalize features
	image_features = image_features / image_features.norm(dim=-1, keepdim=True)
	text_features = text_features / text_features.norm(dim=-1, keepdim=True)

	# Calculate similarities
	similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
	values, indices = similarities[0].topk(top_k)

	# Format results
	results = []
	for i in range(top_k):
	category = CATEGORIES[indices[i]]
	confidence = float(values[i])
	results.append((category, confidence))

	if word.lower() not in [cat.lower() for cat, _ in results]:
	# Find the original confidence score for the word
	word_idx = CATEGORIES.index(word)
	word_confidence = float(similarities[0][word_idx])

	# Remove the last result if we have top_k results
	if len(results) >= top_k:
	results.pop()
	# Add the word with its original confidence score
	results.append((word, word_confidence))

	return results

	except Exception as e:
	print(f"Error in classification: {str(e)}")
	return []

	def gradio_classify(image):
	"""Gradio interface function"""
	if image is None:
	return "Please draw something first!"

	results = classify_drawing(image, top_k=5)

	if not results:
	return "Sorry, couldn't classify the drawing. Please try again."

	# Format results for display
	output = "🎨 Drawing Classification Results:\n\n"
	for i, (category, confidence) in enumerate(results, 1):
	percentage = confidence * 100
	output += f"{i}. {category.title()} - {percentage:.1f}%\n"

	return output

	# Create Gradio interface
	def create_interface():
	"""Create the Gradio interface with both drawing and API tabs"""

	with gr.Blocks(
	title="QuickDraw Classifier",
	theme=gr.themes.Soft(),
	css="""
	.main-header { text-align: center; margin-bottom: 2rem; }
	.description { text-align: center; margin-bottom: 1.5rem; font-size: 1.1rem; }
	"""
	) as demo:

	gr.HTML("""
	<div class="main-header">
	<h1>🎨 QuickDraw Classifier</h1>
	<p class="description">
	Draw anything and let AI guess what it is! This app uses CLIP for zero-shot classification
	of your drawings across 345+ categories from Google's QuickDraw dataset.
	</p>
	</div>
	""")

	with gr.Tabs():
	# Main Drawing Interface Tab
	with gr.Tab("🖊️ Draw & Classify"):
	with gr.Row():
	with gr.Column(scale=1):
	# Drawing canvas
	canvas = gr.Sketchpad(
	label="Draw here! 🖊️",
	type="pil",
	image_mode="RGB",
	canvas_size=(400, 400),
	brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=5),
	show_download_button=False,
	show_share_button=False,
	)

	with gr.Row():
	classify_btn = gr.Button("🔍 Classify Drawing", variant="primary", size="lg")
	clear_btn = gr.ClearButton(canvas, value="🗑️ Clear Canvas", size="lg")

	with gr.Column(scale=1):
	# Results
	output = gr.Markdown(
	label="Classification Results",
	value="👈 Draw something and click 'Classify Drawing' to see results!"
	)

	# Examples section
	gr.HTML("""
	<div style="margin-top: 2rem; text-align: center;">
	<h3>💡 Drawing Tips:</h3>
	<p>• Draw simple, clear shapes<br>
	• Try common objects like animals, vehicles, or everyday items<br>
	• The AI knows 345+ categories including: airplane, cat, house, tree, car, flower, etc.</p>
	</div>
	""")

	# Event handlers
	classify_btn.click(
	fn=gradio_classify,
	inputs=[canvas],
	outputs=[output]
	)

	# Auto-classify on canvas change (optional, can be removed if too frequent)
	canvas.change(
	fn=gradio_classify,
	inputs=[canvas],
	outputs=[output]
	)

	# API Interface Tab
	with gr.Tab("🔌 API Testing"):
	gr.HTML("""
	<div style="text-align: center; margin-bottom: 1rem;">
	<h3>Base64 Image Classification API</h3>
	<p>Test the API endpoint with base64 encoded images</p>
	</div>
	""")

	with gr.Row():
	with gr.Column():
	api_input = gr.Textbox(
	label="Base64 Image Data",
	placeholder="Paste base64 encoded image here...",
	lines=3
	)
	api_top_k = gr.Slider(
	minimum=1,
	maximum=10,
	value=5,
	step=1,
	label="Top K Predictions"
	)
	api_word = gr.Textbox(
	label="Word to Add",
	placeholder="Enter a word to add to the categories",
	lines=1
	)
	api_btn = gr.Button("🚀 Classify", variant="primary")

	with gr.Column():
	api_output = gr.JSON(label="API Response")

	api_btn.click(
	fn=classify_image_api,
	inputs=[api_input, api_top_k, api_word],
	outputs=[api_output]
	)

	return demo

	# API endpoint function for external use
	def classify_image_api(image_data: str, top_k: int = 5, word: str = "dog"):
	"""
	API endpoint for classifying base64 encoded images

	Args:
	image_data: Base64 encoded image string
	top_k: Number of top predictions to return

	Returns:
	JSON response with classification results
	"""
	try:
	# Decode base64 image
	image_bytes = base64.b64decode(image_data)
	image = Image.open(io.BytesIO(image_bytes))

	# Classify
	results = classify_drawing1(image, top_k, word)
	# Normalize confidence scores to sum to 100%
	total_conf = sum(conf for _, conf in results)
	if total_conf > 0:
	results = [(cat, (conf/total_conf) * 100) for cat, conf in results]

	return {
	"success": True,
	"predictions": [
	{"category": cat, "confidence": conf}
	for cat, conf in results
	]
	}

	except Exception as e:
	return {
	"success": False,
	"error": str(e)
	}

	if __name__ == "__main__":
	# Create and launch the main interface
	demo = create_interface()

	demo.launch(
	server_name="0.0.0.0",
	server_port=7860,
	share=True
	)