Create app.py

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import torch
+from torchvision import models, transforms
+from huggingface_hub import hf_hub_download
+
+# --- 1. CONFIGURATION ---
+REPO_ID = "arpit-gour02/document-classification"
+MODEL_FILENAME = "resnet50_epoch_5.pth"
+
+class_names = [
+    'letter', 'form', 'email', 'handwritten', 'advertisement', 'scientific report', 
+    'scientific publication', 'specification', 'file folder', 'news article', 
+    'budget', 'invoice', 'presentation', 'questionnaire', 'resume', 'memo'
+]
+
+# --- 2. LOAD MODEL FROM HUB ---
+def load_model_from_hub():
+    print(f"Downloading {MODEL_FILENAME} from {REPO_ID}...")
+    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
+    
+    # Initialize Architecture
+    model = models.resnet50(num_classes=16)
+    
+    # Load Weights
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+        
+    model.eval()
+    return model
+
+model = load_model_from_hub()
+
+# --- 3. PREPROCESSING ---
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# --- 4. PREDICTION FUNCTION ---
+def predict(image):
+    if image is None:
+        return None
+    image_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+    
+    return {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
+
+# --- 5. LAUNCH INTERFACE ---
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Document Classifier (ResNet50)",
+    description=f"Classifies documents into 16 categories using a ResNet50 model hosted at <a href='https://huggingface.co/{REPO_ID}'>{REPO_ID}</a>.",
+    examples=[["sample_invoice.jpg"]] # You can upload a sample image to the space if you want examples
+)
+
+if __name__ == "__main__":
+    interface.launch()