Upload 6 files

necessary files for deployment

api_app.py

@@ -0,0 +1,62 @@
+from fastapi import FastAPI, File, UploadFile
+from transformers import ViTForImageClassification, AutoImageProcessor
+from transformers import pipeline
+import io
+from PIL import Image
+import torch
+import os
+
+# --- Setup ---
+app = FastAPI()
+# Define the path where the model files will be located in the deployed environment
+MODEL_PATH = "." # In Hugging Face Spaces, files are often in the root directory
+
+# Load model and processor outside the endpoint for speed
+try:
+    processor = AutoImageProcessor.from_pretrained(MODEL_PATH)
+    model = ViTForImageClassification.from_pretrained(MODEL_PATH)
+except Exception as e:
+    # A fallback/debug print if loading fails during deployment
+    print(f"Error loading model or processor from {MODEL_PATH}: {e}")
+    # You might need to check file names or adjust MODEL_PATH if this fails
+    
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
+
+# --- Prediction Endpoint ---
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    """Accepts an image file and returns the probability of a leak."""
+    try:
+        # 1. Read the uploaded file bytes
+        data = await file.read()
+        
+        # 2. Open the image using PIL
+        img = Image.open(io.BytesIO(data)).convert("RGB")
+        
+        # 3. Preprocess the image (resize, normalize)
+        inputs = processor(images=img, return_tensors="pt")
+        inputs = {k:v.to(device) for k,v in inputs.items()}
+        
+        # 4. Run inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            # Apply softmax to get probabilities
+            probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
+            
+            # The 'leak' label has ID 1 (based on your id2label = {0: "no_leak", 1: "leak"})
+            prob_leak = float(probs[1]) 
+            
+        return {
+            "prediction": "leak" if prob_leak >= 0.5 else "no_leak",
+            "leak_probability": prob_leak
+        }
+        
+    except Exception as e:
+        return {"error": str(e), "message": "Prediction failed."}
+
+# --- Root Endpoint (for health check) ---
+@app.get("/")
+def home():
+    return {"status": "Model API is running"}

app.py

@@ -0,0 +1,13 @@
+import os
+from subprocess import Popen
+
+# This script tells the Hugging Face environment to run your FastAPI app
+# on the correct port, which is necessary when using the Gradio SDK template.
+
+# Set the port FastAPI should listen on (HF requires 7860)
+os.environ['PORT'] = '7860'
+
+# Start the uvicorn server to host the FastAPI application
+# api_app:app refers to the 'app' variable inside your 'api_app.py' file
+# --host 0.0.0.0 is necessary for external access
+Popen(["uvicorn", "api_app:app", "--host", "0.0.0.0", "--port", "7860"])

config.json

@@ -0,0 +1,33 @@
+{
+  "architectures": [
+    "ViTForImageClassification"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "dtype": "float32",
+  "encoder_stride": 16,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "no_leak",
+    "1": "leak"
+  },
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "leak": 1,
+    "no_leak": 0
+  },
+  "layer_norm_eps": 1e-12,
+  "model_type": "vit",
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "patch_size": 16,
+  "pooler_act": "tanh",
+  "pooler_output_size": 768,
+  "problem_type": "single_label_classification",
+  "qkv_bias": true,
+  "transformers_version": "4.57.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00f8294b61d9e6fb879e401172c44b5fa2cbfff6a906a1d72a3a5e130725a977
+size 343223968

preprocessor_config.json

@@ -0,0 +1,33 @@
+{
+  "crop_size": null,
+  "data_format": "channels_first",
+  "default_to_square": true,
+  "device": null,
+  "disable_grouping": null,
+  "do_center_crop": null,
+  "do_convert_rgb": null,
+  "do_normalize": true,
+  "do_pad": null,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "ViTImageProcessorFast",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "input_data_format": null,
+  "pad_size": null,
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "return_tensors": null,
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+transformers
+datasets
+torch
+pillow
+scikit-learn