tchdhry/actor-classifier

1import gradio as gr2import torch3from torchvision import models, transforms4from PIL import Image5import os6 7# Setup8device = torch.device('cpu')9print("Loading model...")10 11# Create model architecture12model = models.resnet18(weights=None)13model.fc = torch.nn.Linear(512, 2)14 15# Load weights16try:17    if os.path.exists('actor_model_pytorch.pth'):18        checkpoint = torch.load('actor_model_pytorch.pth', map_location=device)19    else:20        checkpoint = torch.load('best_actor_model.pth', map_location=device)21    22    model.load_state_dict(checkpoint['model_state_dict'])23    model.eval()24    print("Model loaded successfully!")25except Exception as e:26    print(f"Error loading model: {e}")27 28# Simple transform29transform = transforms.Compose([30    transforms.Resize((224, 224)),31    transforms.ToTensor(),32])33 34def predict(image):35    if image is None:36        return "Please upload an image"37    38    try:39        # Convert and predict40        img = transform(image).unsqueeze(0)41        42        with torch.no_grad():43            outputs = model(img)44            probs = torch.nn.functional.softmax(outputs, dim=1)45        46        # Get result47        actor_prob = float(probs[0][0])48        49        # Simple text response50        if actor_prob > 0.7:51            return f"🎬 Definitely an actor! ({actor_prob*100:.1f}% confident)"52        elif actor_prob > 0.5:53            return f"🤔 Probably an actor ({actor_prob*100:.1f}% confident)"54        else:55            return f"👤 Not an actor ({(1-actor_prob)*100:.1f}% confident)"56            57    except Exception as e:58        return f"Error: {str(e)}"59 60# Create the simplest possible interface61demo = gr.Interface(62    fn=predict,63    inputs=gr.Image(type="pil"),64    outputs=gr.Textbox(label="Result"),65    title="Actor or Not? 🎭",66    description="Upload a photo to check if someone looks like an actor. (Learning project - trained on limited data!)"67)68 69if __name__ == "__main__":70    demo.launch()71 72