Spaces:

Shad0wKillar
/

EfficientNet_TransferLearned

Sleeping

App Files Files Community

Shad0wKillar commited on 21 days ago

Commit

4c4efcd

verified ·

1 Parent(s): f9d30f0

Added multi-modal support.

Browse files

Files changed (1) hide show

main.py +80 -189

main.py CHANGED Viewed

@@ -1,5 +1,5 @@
 import io
-from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import HTMLResponse
 import torch
 import torchvision
@@ -9,224 +9,112 @@ 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.load_state_dict(
-    torch.load(weights_path, map_location=torch.device("cpu"), weights_only=True)
-)
-model.eval()
-transform = torchvision.transforms.Compose(
-    [
-        torchvision.transforms.Resize(255, interpolation=InterpolationMode.BILINEAR),
-        torchvision.transforms.CenterCrop(240),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize(
-            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-        ),
-    ]
-)
 class_names = ["pizza", "steak", "sushi"]
 @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]);
                 try {
-                    const response = await fetch("/predict", {
-                        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>
     </body>
@@ -234,16 +122,19 @@ async def read_root():
     """
     return HTMLResponse(content=html_content)
 @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))}

 import io
+from fastapi import FastAPI, UploadFile, File, Query
 from fastapi.responses import HTMLResponse
 import torch
 import torchvision
 app = FastAPI()
+# Map model names to their specific configuration and repo details
+MODEL_CONFIGS = {
+    "b1": {"repo": "Shad0wKillar/efficientnet-b1", "file": "EfficientNet_B1_20percent.pth", "features": 1280},
+    "b3": {"repo": "Shad0wKillar/efficientnet-b3", "file": "EfficientNet_B3_20percent.pth", "features": 1536},
+    "b5": {"repo": "Shad0wKillar/efficientnet-b5", "file": "EfficientNet_B5_20percent.pth", "features": 2048},
+    "b7": {"repo": "Shad0wKillar/efficientnet-b7", "file": "EfficientNet_B7_20percent.pth", "features": 2560},
+}
+def create_model(model_type):
+    # I matched the model architecture to the specific version
+    if model_type == "b1": model = torchvision.models.efficientnet_b1()
+    elif model_type == "b3": model = torchvision.models.efficientnet_b3()
+    elif model_type == "b5": model = torchvision.models.efficientnet_b5()
+    elif model_type == "b7": model = torchvision.models.efficientnet_b7()
     model.classifier = torch.nn.Sequential(
         torch.nn.Dropout(p=0.2, inplace=True),
+        torch.nn.Linear(in_features=MODEL_CONFIGS[model_type]["features"], out_features=3, bias=True),
     )
     return model
+# I pre-loaded all models into memory for fast access
+loaded_models = {}
+for m_type, config in MODEL_CONFIGS.items():
+    m = create_model(m_type)
+    path = hf_hub_download(repo_id=config["repo"], filename=config["file"])
+    m.load_state_dict(torch.load(path, map_location=torch.device("cpu"), weights_only=True))
+    m.eval()
+    loaded_models[m_type] = m
+transform = torchvision.transforms.Compose([
+    torchvision.transforms.Resize(255, interpolation=InterpolationMode.BILINEAR),
+    torchvision.transforms.CenterCrop(240),
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
 class_names = ["pizza", "steak", "sushi"]
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
+    # I added a dropdown to the UI to select the model
     html_content = """
     <!DOCTYPE html>
     <html lang="en">
     <head>
         <meta charset="UTF-8">
         <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>EfficientNet Multi-Model API</title>
         <style>
+            body { font-family: system-ui, 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; }
+            select, input[type="file"], button { width: 100%; padding: 10px; margin-bottom: 15px; border-radius: 6px; border: 1px solid #374151; background: #0b0f19; color: #e5e7eb; box-sizing: border-box; }
+            button { background: #ffffff; color: #000; font-weight: 600; cursor: pointer; border: none; }
+            #preview { max-width: 100%; border-radius: 6px; display: none; margin-bottom: 15px; }
+            .result-box { padding: 15px; background: #0b0f19; border-radius: 6px; display: none; text-align: center; }
         </style>
     </head>
     <body>
         <div class="container">
+            <h2>Food Classification</h2>
+            <label for="modelSelect">Select Model Architecture:</label>
+            <select id="modelSelect">
+                <option value="b1">EfficientNet-B1</option>
+                <option value="b3">EfficientNet-B3</option>
+                <option value="b5">EfficientNet-B5</option>
+                <option value="b7">EfficientNet-B7</option>
+            </select>
+            <input type="file" id="imageInput" accept="image/*" onchange="previewImage(event)">
+            <img id="preview">
             <button onclick="testAPI()" id="runBtn">Run Prediction</button>
             <div class="result-box" id="resultBox">
+                <div id="topPrediction" style="font-size: 1.5rem; color: #10b981; font-weight: 700;"></div>
+                <div id="rawProbs" style="font-size: 0.75rem; color: #6b7280; margin-top: 10px;"></div>
             </div>
         </div>
         <script>
             function previewImage(event) {
                 const reader = new FileReader();
+                reader.onload = () => {
+                    const p = document.getElementById('preview');
+                    p.src = reader.result; p.style.display = 'block';
                 };
+                reader.readAsDataURL(event.target.files[0]);
             }
             async function testAPI() {
+                const file = document.getElementById('imageInput').files[0];
+                const model = document.getElementById('modelSelect').value;
+                if (!file) return alert("Select an image");
+                const btn = document.getElementById('runBtn');
+                btn.innerText = "Processing..."; btn.disabled = true;
                 const formData = new FormData();
+                formData.append("file", file);
                 try {
+                    const res = await fetch(`/predict?model_type=${model}`, { method: "POST", body: formData });
+                    const data = await res.json();
+                    const best = Object.entries(data).reduce((a, b) => a[1] > b[1] ? a : b);
+                    document.getElementById('topPrediction').innerText = best[0].toUpperCase();
+                    document.getElementById('rawProbs').innerText = JSON.stringify(data);
+                    document.getElementById('resultBox').style.display = 'block';
+                } catch (e) { alert("Error: " + e.message); }
+                finally { btn.innerText = "Run Prediction"; btn.disabled = false; }
             }
         </script>
     </body>
     """
     return HTMLResponse(content=html_content)
 @app.post("/predict")
+async def predict(model_type: str = Query("b1"), file: UploadFile = File(...)):
+    # I routed the request to the specific loaded model
+    if model_type not in loaded_models:
+        return {"error": "Model not found"}
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     img_tensor = transform(image).unsqueeze(0)
+    selected_model = loaded_models[model_type]
     with torch.no_grad():
+        logits = selected_model(img_tensor)
         probs = torch.softmax(logits, dim=1).squeeze()
     return {class_names[i]: float(probs[i]) for i in range(len(class_names))}