Rename app.py to main.py

app.py

@@ -1,7 +0,0 @@
-# app.py
-from main import app
-
-# For Hugging Face Spaces
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)

main.py

@@ -0,0 +1,463 @@
+from fasthtml.common import *
+from fastai.vision.all import *
+import os
+import time
+from pathlib import Path
+import urllib.request
+from io import BytesIO
+
+# Create necessary directories
+os.makedirs('uploads', exist_ok=True)
+
+# Function to load model - with fallback for testing
+def load_model():
+    try:
+        model_path = 'levit.pkl'
+        # Check if model exists, if not try to download a sample model (for demo purposes)
+        if not os.path.exists(model_path):
+            print("Model not found. This is just for testing purposes.")
+            # In a real deployment, you'd want to handle this more gracefully
+            return None, ['class1', 'class2', 'class3']
+        
+        learn = load_learner(model_path)
+        labels = learn.dls.vocab
+        print(f"Model loaded successfully with labels: {labels}")
+        return learn, labels
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        # Fallback for testing
+        return None, ['class1', 'class2', 'class3']
+
+# Load the model at startup
+learn, labels = load_model()
+
+# Create a FastHTML app
+app, rt = fast_app()
+
+# Define the prediction function
+def predict(img_bytes):
+    try:
+        # If no model is loaded, return mock predictions for testing
+        if learn is None:
+            import random
+            mock_results = {label: random.random() for label in labels}
+            # Sort by values and normalize to ensure they sum to 1
+            total = sum(mock_results.values())
+            return {k: v/total for k, v in sorted(mock_results.items(), key=lambda x: x[1], reverse=True)}
+        
+        # Real prediction with the model
+        img = PILImage.create(BytesIO(img_bytes))
+        img = img.resize((512, 512))
+        pred, pred_idx, probs = learn.predict(img)
+        return {labels[i]: float(probs[i]) for i in range(len(labels))}
+    except Exception as e:
+        print(f"Prediction error: {e}")
+        return {"Error": 1.0}
+
+# Main page route
+@rt("/")
+def get():
+    # Create a form for image upload
+    upload_form = Form(
+        Div(
+            H1("FastAI Image Classifier"),
+            P("Upload an image to classify it using a pre-trained model."),
+            cls="instructions"
+        ),
+        Div(
+            Input(type="file", name="image", accept="image/*", required=True, 
+                  hx_indicator="#loading"),
+            Button("Classify", type="submit"),
+            cls="upload-controls"
+        ),
+        hx_post="/predict",
+        hx_target="#result",
+        hx_swap="innerHTML",
+        hx_encoding="multipart/form-data",
+        id="upload-form"
+    )
+    
+    # Add loading indicator
+    loading = Div(
+        P("Processing your image..."),
+        id="loading",
+        cls="htmx-indicator"
+    )
+    
+    # Container for results
+    result_container = Div(id="result", cls="result-container")
+    
+    # Example section
+    examples = Div(
+        H2("Or try an example:"),
+        A("Example Image", href="#", 
+          hx_get="/predict_example", 
+          hx_target="#result",
+          hx_indicator="#loading"),
+        cls="examples-section"
+    )
+    
+    # CSS styles
+    css = """
+    :root {
+        --primary-color: #3498db;
+        --secondary-color: #2c3e50;
+        --background-color: #f9f9f9;
+        --error-color: #e74c3c;
+        --shadow-color: rgba(0, 0, 0, 0.1);
+        --border-color: #ddd;
+    }
+    
+    body {
+        font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', 'Helvetica Neue', sans-serif;
+        line-height: 1.6;
+        color: #333;
+        max-width: 800px;
+        margin: 0 auto;
+        padding: 20px;
+        background-color: #fff;
+    }
+    
+    h1 {
+        color: var(--secondary-color);
+        margin-bottom: 1rem;
+        font-weight: 600;
+    }
+    
+    h2 {
+        color: var(--primary-color);
+        margin-top: 1.5rem;
+        font-weight: 500;
+    }
+    
+    .instructions {
+        margin-bottom: 20px;
+    }
+    
+    .upload-controls {
+        display: flex;
+        gap: 10px;
+        margin-bottom: 30px;
+        align-items: center;
+        flex-wrap: wrap;
+    }
+    
+    button {
+        background-color: var(--primary-color);
+        color: white;
+        border: none;
+        padding: 10px 15px;
+        border-radius: 4px;
+        cursor: pointer;
+        transition: background-color 0.3s;
+        font-weight: 500;
+    }
+    
+    button:hover {
+        background-color: #2980b9;
+    }
+    
+    input[type="file"] {
+        padding: 10px;
+        border: 1px solid var(--border-color);
+        border-radius: 4px;
+        flex-grow: 1;
+    }
+    
+    #upload-form {
+        margin-bottom: 40px;
+        padding: 20px;
+        border-radius: 8px;
+        background-color: var(--background-color);
+        box-shadow: 0 2px 10px var(--shadow-color);
+    }
+    
+    .result-container {
+        margin-top: 20px;
+    }
+    
+    .prediction-results {
+        margin-top: 20px;
+        padding: 20px;
+        border: 1px solid var(--border-color);
+        border-radius: 8px;
+        background-color: var(--background-color);
+        box-shadow: 0 2px 8px var(--shadow-color);
+    }
+    
+    .result-image {
+        max-width: 100%;
+        height: auto;
+        border-radius: 8px;
+        box-shadow: 0 2px 5px var(--shadow-color);
+        margin-bottom: 20px;
+        display: block;
+    }
+    
+    .prediction-list {
+        margin-top: 15px;
+    }
+    
+    .prediction-item {
+        padding: 12px 15px;
+        margin-bottom: 10px;
+        background-color: white;
+        border-radius: 6px;
+        box-shadow: 0 1px 3px var(--shadow-color);
+    }
+    
+    .label-text {
+        margin-bottom: 8px;
+        font-weight: 500;
+        display: flex;
+        justify-content: space-between;
+    }
+    
+    .examples-section {
+        margin-top: 30px;
+        padding-top: 20px;
+        border-top: 1px solid var(--border-color);
+    }
+    
+    .htmx-indicator {
+        display: none;
+        padding: 15px;
+        background-color: #e8f4fc;
+        border-radius: 6px;
+        text-align: center;
+        margin: 15px 0;
+        box-shadow: 0 1px 3px var(--shadow-color);
+    }
+    
+    .htmx-request .htmx-indicator {
+        display: block;
+    }
+    
+    .progress-bar {
+        height: 10px;
+        background-color: #f0f0f0;
+        border-radius: 5px;
+        margin: 5px 0;
+        overflow: hidden;
+    }
+    
+    .progress-fill {
+        height: 100%;
+        background-color: var(--primary-color);
+        width: 0;
+        transition: width 0.5s ease;
+    }
+    
+    .error-message {
+        color: var(--error-color);
+        padding: 15px;
+        border: 1px solid var(--error-color);
+        border-radius: 5px;
+        background-color: #fde9e7;
+    }
+    
+    a {
+        color: var(--primary-color);
+        text-decoration: none;
+        font-weight: 500;
+    }
+    
+    a:hover {
+        text-decoration: underline;
+    }
+    
+    /* Responsive styling */
+    @media (max-width: 600px) {
+        .upload-controls {
+            flex-direction: column;
+            align-items: stretch;
+        }
+        
+        button {
+            width: 100%;
+        }
+    }
+    
+    .model-info {
+        font-size: 0.9rem;
+        color: #666;
+        margin-top: 40px;
+        padding-top: 20px;
+        border-top: 1px solid var(--border-color);
+    }
+    """
+    
+    # Model information
+    model_info = Div(
+        P(f"Model: {'Model loaded successfully' if learn is not None else 'Demo mode - no model loaded'}"),
+        P(f"Classes: {', '.join(labels)}"),
+        cls="model-info"
+    )
+    
+    return Titled("FastAI Image Classifier", 
+                 upload_form,
+                 loading,
+                 result_container,
+                 examples,
+                 model_info,
+                 Style(css))
+
+# Prediction route for uploaded images
+@rt("/predict")
+async def post(image: UploadFile):
+    try:
+        # Read the uploaded image
+        image_bytes = await image.read()
+        
+        # Generate a unique filename to avoid conflicts
+        from datetime import datetime
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        safe_filename = f"{timestamp}_{image.filename.replace(' ', '_')}"
+        
+        # Save the image temporarily
+        img_path = f"uploads/{safe_filename}"
+        with open(img_path, "wb") as f:
+            f.write(image_bytes)
+        
+        # Add a small delay to make the loading indicator visible
+        time.sleep(0.5)
+        
+        # Make a prediction
+        results = predict(image_bytes)
+        
+        # Sort results by probability
+        sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+        top_results = dict(list(sorted_results.items())[:3])
+        
+        # Create prediction items with progress bars
+        prediction_items = []
+        for label, prob in top_results.items():
+            percentage = int(prob * 100)
+            prediction_items.append(
+                Div(
+                    Div(
+                        Span(f"{label}"),
+                        Span(f"{percentage}%"),
+                        cls="label-text"
+                    ),
+                    Div(
+                        Div(cls="progress-fill", style=f"width: {percentage}%;"),
+                        cls="progress-bar"
+                    ),
+                    cls="prediction-item"
+                )
+            )
+        
+        # Create result HTML
+        result_html = Div(
+            H2("Prediction Results:"),
+            Img(src=f"/image/{safe_filename}", cls="result-image", alt="Uploaded image"),
+            Div(*prediction_items, cls="prediction-list"),
+            cls="prediction-results"
+        )
+        
+        return result_html
+    
+    except Exception as e:
+        return Div(
+            H2("Error"),
+            P(f"An error occurred during prediction: {str(e)}"),
+            cls="error-message"
+        )
+
+# Route to serve saved images
+@rt("/image/{filename}")
+def get(filename: str):
+    file_path = f"uploads/{filename}"
+    if os.path.exists(file_path):
+        return FileResponse(file_path)
+    else:
+        return Div(
+            H2("Error"),
+            P("Image not found."),
+            cls="error-message"
+        )
+
+# Route for example image
+@rt("/predict_example")
+def get():
+    try:
+        # Path to example image
+        example_path = "image.jpg"
+        
+        # Check if example image exists
+        if os.path.exists(example_path):
+            with open(example_path, "rb") as f:
+                image_bytes = f.read()
+            
+            # Save the example image to uploads
+            example_name = "example.jpg"
+            with open(f"uploads/{example_name}", "wb") as f:
+                f.write(image_bytes)
+            
+            # Add a small delay to make the loading indicator visible
+            time.sleep(0.5)
+            
+            # Make a prediction
+            results = predict(image_bytes)
+            
+            # Sort results by probability
+            sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+            top_results = dict(list(sorted_results.items())[:3])
+            
+            # Create prediction items with progress bars
+            prediction_items = []
+            for label, prob in top_results.items():
+                percentage = int(prob * 100)
+                prediction_items.append(
+                    Div(
+                        Div(
+                            Span(f"{label}"),
+                            Span(f"{percentage}%"),
+                            cls="label-text"
+                        ),
+                        Div(
+                            Div(cls="progress-fill", style=f"width: {percentage}%;"),
+                            cls="progress-bar"
+                        ),
+                        cls="prediction-item"
+                    )
+                )
+            
+            # Create result HTML
+            result_html = Div(
+                H2("Prediction Results:"),
+                Img(src=f"/image/{example_name}", cls="result-image", alt="Example image"),
+                Div(*prediction_items, cls="prediction-list"),
+                P("This is a demonstration using the provided example image.", style="font-style: italic; color: #666;"),
+                cls="prediction-results"
+            )
+            
+            return result_html
+        else:
+            return Div(
+                H2("Example Not Found"),
+                P("The example image 'image.jpg' was not found. Please try uploading your own image."),
+                cls="error-message"
+            )
+    
+    except Exception as e:
+        return Div(
+            H2("Error"),
+            P(f"An error occurred with the example: {str(e)}"),
+            cls="error-message"
+        )
+
+# Health check endpoint (useful for Docker/Kubernetes)
+@rt("/health")
+def get():
+    return {"status": "ok", "model_loaded": learn is not None}
+
+# Run the app
+if __name__ == "__main__":
+    # Use environment variables if available (common in Docker)
+    host = os.environ.get("HOST", "0.0.0.0")
+    port = int(os.environ.get("PORT", 8000))
+    
+    print(f"Starting FastHTML server on {host}:{port}")
+    serve(app=app, host=host, port=port)