Add image saving logic secure

app/main.py

@@ -7,6 +7,10 @@ from app.inference import load_classification_model, classify_bytes
 from app.inference import load_classification_model, classify_bytes
 from app.inference_yolo import classify_yolo_bytes, load_yolo_model
 # from app.model import load_model, predict_pca_from_bytes
+import io  
+import os
+import uuid
+from huggingface_hub import HfApi
 import json, os
 
 
@@ -33,6 +37,37 @@ INDEX_HTML = os.path.join(STATIC_DIR, "index.html")
 
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
+
+# --- CONFIGURATION ---
+HF_TOKEN = os.environ.get("HF_TOKEN") 
+
+DATASET_REPO_ID = "AndrewKof/NEMO-user-uploads" 
+
+api = HfApi(token=HF_TOKEN)
+
+def save_image_to_hub(image_bytes):
+    """
+    Uploads raw image bytes to the dataset.
+    """
+    # Create a unique filename 
+    filename = f"user_images/{uuid.uuid4()}.png"
+    
+    try:
+        print(f"Uploading {filename}...")
+        
+        # Wrap bytes in BytesIO so it looks like a file
+        file_object = io.BytesIO(image_bytes) 
+        
+        api.upload_file(
+            path_or_fileobj=file_object, 
+            path_in_repo=filename,
+            repo_id=DATASET_REPO_ID,
+            repo_type="dataset"
+        )
+        print("Upload successful!")
+    except Exception as e:
+        print(f"Error uploading image: {e}")
+
 @app.get("/", response_class=HTMLResponse)
 def serve_frontend():
     """Serve the web interface."""
@@ -64,6 +99,7 @@ print("✅ Classification model loaded and ready for inference!")
 async def generate_attention(file: UploadFile = File(...)):
     """Generate and return mean attention map for uploaded image."""
     image_bytes = await file.read()
+    save_image_to_hub(image_bytes)
     result = predict_from_bytes(model_device_tuple, image_bytes)
     return result
 
@@ -73,7 +109,7 @@ async def classify(
     model: str = Form("dino")  # <--- Read 'model' from FormData (default 'dino')
 ):
     image_bytes = await file.read()
-    
+    save_image_to_hub(image_bytes)
     if model == "yolo":
         print("🧠 Running YOLOv11 Inference...")
         return classify_yolo_bytes(image_bytes)