Spaces:

Shad0wKillar
/

EfficientNet_TransferLearned

Sleeping

App Files Files Community

Shad0wKillar commited on 27 days ago

Commit

f9d30f0

verified ·

1 Parent(s): c6403bc

Added better UI

Browse files

Files changed (1) hide show

main.py +156 -30

main.py CHANGED Viewed

@@ -1,28 +1,23 @@
 import io
 from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import HTMLResponse  # Add this import
 import torch
 import torchvision
 from torchvision.transforms import InterpolationMode
-from fastapi import FastAPI, UploadFile, File
 from huggingface_hub import hf_hub_download
 from PIL import Image
 app = FastAPI()
-# I recreated the exact architecture from your training script
 def create_model():
     model = torchvision.models.efficientnet_b1()
-    # I replaced the classifier exactly as you did in training
     model.classifier = torch.nn.Sequential(
         torch.nn.Dropout(p=0.2, inplace=True),
         torch.nn.Linear(in_features=1280, out_features=3, bias=True),
     )
     return model
-# I load the model and weights when the Docker container starts
 model = create_model()
 weights_path = hf_hub_download(
     repo_id="Shad0wKillar/efficientnet-b1", filename="EfficientNet_B1_20percent.pth"
@@ -32,7 +27,6 @@ model.load_state_dict(
 )
 model.eval()
-# I mapped the exact auto_transform sequence you provided
 transform = torchvision.transforms.Compose(
     [
         torchvision.transforms.Resize(255, interpolation=InterpolationMode.BILINEAR),
@@ -49,41 +43,152 @@ class_names = ["pizza", "steak", "sushi"]
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
     html_content = """
     <!DOCTYPE html>
-    <html>
     <head>
         <title>Model Testing API</title>
         <style>
-            body { font-family: sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; background: #121212; color: #ffffff; }
-            .box { border: 1px solid #333; padding: 20px; border-radius: 8px; background: #1e1e1e; }
-            button { background: #3b82f6; color: white; border: none; padding: 10px 15px; border-radius: 4px; cursor: pointer; margin-top: 10px; }
-            button:hover { background: #2563eb; }
-            pre { background: #000; padding: 15px; border-radius: 4px; overflow-x: auto; color: #10b981; }
         </style>
     </head>
     <body>
-        <div class="box">
-            <h2>Backend Diagnostic UI</h2>
-            <p>Upload an image to test the API directly.</p>
-            <input type="file" id="imageInput" accept="image/jpeg, image/png">
-            <button onclick="testAPI()">Run Prediction</button>
-            <h3 style="margin-top: 20px;">Response:</h3>
-            <pre id="output">Awaiting image...</pre>
         </div>
         <script>
             async function testAPI() {
                 const fileInput = document.getElementById('imageInput');
-                const output = document.getElementById('output');
                 if (fileInput.files.length === 0) {
-                    output.innerText = "Please select an image first.";
                     return;
                 }
-                output.innerText = "Processing...";
                 const formData = new FormData();
                 formData.append("file", fileInput.files[0]);
@@ -93,10 +198,34 @@ async def read_root():
                         method: "POST",
                         body: formData
                     });
                     const data = await response.json();
-                    output.innerText = JSON.stringify(data, null, 2);
                 } catch (error) {
-                    output.innerText = "Error hitting API: " + error.message;
                 }
             }
         </script>
@@ -108,16 +237,13 @@ async def read_root():
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # I read the incoming bytes into a PIL image
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    # I process the image and run inference
     img_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         logits = model(img_tensor)
         probs = torch.softmax(logits, dim=1).squeeze()
-    # I format the output as a JSON dictionary
-    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}

 import io
 from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import HTMLResponse
 import torch
 import torchvision
 from torchvision.transforms import InterpolationMode
 from huggingface_hub import hf_hub_download
 from PIL import Image
 app = FastAPI()
 def create_model():
     model = torchvision.models.efficientnet_b1()
     model.classifier = torch.nn.Sequential(
         torch.nn.Dropout(p=0.2, inplace=True),
         torch.nn.Linear(in_features=1280, out_features=3, bias=True),
     )
     return model
+# I loaded the model and weights for the container
 model = create_model()
 weights_path = hf_hub_download(
     repo_id="Shad0wKillar/efficientnet-b1", filename="EfficientNet_B1_20percent.pth"
 )
 model.eval()
 transform = torchvision.transforms.Compose(
     [
         torchvision.transforms.Resize(255, interpolation=InterpolationMode.BILINEAR),
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
+    # I built the new UI matching the Hugging Face dark theme
     html_content = """
     <!DOCTYPE html>
+    <html lang="en">
     <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
         <title>Model Testing API</title>
         <style>
+            body {
+                font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+                background-color: #0b0f19;
+                color: #e5e7eb;
+                display: flex;
+                justify-content: center;
+                align-items: center;
+                min-height: 100vh;
+                margin: 0;
+                padding: 20px;
+            }
+            .container {
+                background-color: #1e293b;
+                border: 1px solid #374151;
+                border-radius: 8px;
+                padding: 30px;
+                width: 100%;
+                max-width: 450px;
+                box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
+            }
+            h2 {
+                margin-top: 0;
+                font-size: 1.5rem;
+                font-weight: 600;
+                color: #f3f4f6;
+            }
+            .subtitle {
+                color: #9ca3af;
+                font-size: 0.875rem;
+                margin-bottom: 20px;
+            }
+            input[type="file"] {
+                display: block;
+                width: 100%;
+                padding: 10px;
+                font-size: 0.875rem;
+                color: #9ca3af;
+                background-color: #0b0f19;
+                border: 1px solid #374151;
+                border-radius: 6px;
+                cursor: pointer;
+                box-sizing: border-box;
+                margin-bottom: 15px;
+            }
+            #preview {
+                max-width: 100%;
+                max-height: 300px;
+                object-fit: cover;
+                border-radius: 6px;
+                display: none;
+                margin-bottom: 15px;
+                border: 1px solid #374151;
+            }
+            button {
+                background-color: #ffffff;
+                color: #000000;
+                font-weight: 600;
+                padding: 10px 15px;
+                border: none;
+                border-radius: 6px;
+                cursor: pointer;
+                width: 100%;
+                transition: background-color 0.2s;
+            }
+            button:hover { background-color: #e5e7eb; }
+            button:disabled { background-color: #9ca3af; cursor: not-allowed; }
+            .result-box {
+                margin-top: 20px;
+                padding: 15px;
+                background-color: #0b0f19;
+                border: 1px solid #374151;
+                border-radius: 6px;
+                display: none;
+                text-align: center;
+            }
+            .prediction {
+                font-size: 1.5rem;
+                font-weight: 700;
+                color: #10b981;
+                margin-bottom: 8px;
+            }
+            .raw-probs {
+                font-size: 0.75rem;
+                color: #6b7280;
+                margin: 0;
+            }
         </style>
     </head>
     <body>
+        <div class="container">
+            <h2>Image Classification</h2>
+            <div class="subtitle">Upload an image to test the API endpoint.</div>
+            <input type="file" id="imageInput" accept="image/jpeg, image/png" onchange="previewImage(event)">
+            <img id="preview" alt="Image preview">
+            <button onclick="testAPI()" id="runBtn">Run Prediction</button>
+            <div class="result-box" id="resultBox">
+                <div class="prediction" id="topPrediction"></div>
+                <div class="raw-probs" id="rawProbs"></div>
+            </div>
         </div>
         <script>
+            function previewImage(event) {
+                const reader = new FileReader();
+                reader.onload = function(){
+                    const preview = document.getElementById('preview');
+                    preview.src = reader.result;
+                    preview.style.display = 'block';
+                    document.getElementById('resultBox').style.display = 'none';
+                };
+                if (event.target.files[0]) {
+                    reader.readAsDataURL(event.target.files[0]);
+                }
+            }
             async function testAPI() {
                 const fileInput = document.getElementById('imageInput');
+                const resultBox = document.getElementById('resultBox');
+                const topPrediction = document.getElementById('topPrediction');
+                const rawProbs = document.getElementById('rawProbs');
+                const runBtn = document.getElementById('runBtn');
                 if (fileInput.files.length === 0) {
+                    alert("Please select an image first.");
                     return;
                 }
+                runBtn.innerText = "Processing...";
+                runBtn.disabled = true;
+                resultBox.style.display = 'block';
+                topPrediction.innerText = "Analyzing...";
+                topPrediction.style.color = "#e5e7eb";
+                rawProbs.innerText = "";
                 const formData = new FormData();
                 formData.append("file", fileInput.files[0]);
                         method: "POST",
                         body: formData
                     });
+                    if (!response.ok) throw new Error("API request failed");
                     const data = await response.json();
+                    let highestClass = "";
+                    let highestProb = -1;
+                    let probsList = [];
+                    for (const [className, prob] of Object.entries(data)) {
+                        if (prob > highestProb) {
+                            highestProb = prob;
+                            highestClass = className;
+                        }
+                        probsList.push(`${className}: ${(prob * 100).toFixed(1)}%`);
+                    }
+                    topPrediction.style.color = "#10b981";
+                    topPrediction.innerText = highestClass.charAt(0).toUpperCase() + highestClass.slice(1);
+                    rawProbs.innerText = probsList.join(" • ");
                 } catch (error) {
+                    topPrediction.style.color = "#ef4444";
+                    topPrediction.innerText = "Error";
+                    rawProbs.innerText = error.message;
+                } finally {
+                    runBtn.innerText = "Run Prediction";
+                    runBtn.disabled = false;
                 }
             }
         </script>
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     img_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         logits = model(img_tensor)
         probs = torch.softmax(logits, dim=1).squeeze()
+    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}