Upload app.py

app.py

@@ -0,0 +1,47 @@
+from huggingface_hub import hf_hub_download
+from transformers import ViTForImageClassification, AutoImageProcessor
+
+
+
+from fastapi import FastAPI, File, UploadFile
+from transformers import ViTForImageClassification, AutoImageProcessor
+import io
+from PIL import Image
+import torch
+
+model_file = hf_hub_download(
+    repo_id="iwin10s/leak-detection-model",
+    filename="model.safetensors"
+)
+
+processor_config = hf_hub_download(
+    repo_id="iwin10s/leak-detection-model",
+    filename="preprocessor_config.json"
+)
+
+# Load model & processor
+processor = AutoImageProcessor.from_pretrained("iwin10s/leak-detection-model")
+model = ViTForImageClassification.from_pretrained(
+    pretrained_model_name_or_path="iwin10s/leak-detection-model",
+    local_files_only=False
+)
+
+app = FastAPI()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    data = await file.read()
+
+    img = Image.open(io.BytesIO(data)).convert("RGB")
+    inputs = processor(images=img, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
+
+    return {"leak_probability": float(probs[1])}