Update main.py

main.py

@@ -1,5 +1,3 @@
-# main.py
-
 import base64
 import io
 import json
@@ -16,6 +14,7 @@ import timm
 import segmentation_models_pytorch as smp
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
+import os
 
 # --- 1. SETUP: Create the FastAPI app ---
 app = FastAPI(title="AI Skin Lesion Analyzer API")
@@ -38,13 +37,20 @@ class ImageRequest(BaseModel):
 def load_assets():
     """Load all models and assets from Hugging Face Hub when the server starts."""
     global segmentation_model, classification_model, knowledge_base, transform_segment, transform_classify
-    
+
     print("--> API starting up: This may take a few minutes...")
-    
+
+    # ✅ Create safe folder to store downloaded models
+    os.makedirs("./models_cache", exist_ok=True)
+
     # Load Segmentation Model
     try:
         print("    Downloading UNet segmentation model...")
-        seg_model_path = hf_hub_download(repo_id="sheikh987/unet-isic2018", filename="unet_full_data_best_model.pth")
+        seg_model_path = hf_hub_download(
+            repo_id="sheikh987/unet-isic2018",
+            filename="unet_full_data_best_model.pth",
+            cache_dir="./models_cache"
+        )
         segmentation_model = smp.Unet("resnet34", encoder_weights=None, in_channels=3, classes=1).to(DEVICE)
         segmentation_model.load_state_dict(torch.load(seg_model_path, map_location=DEVICE))
         segmentation_model.eval()
@@ -55,7 +61,11 @@ def load_assets():
     # Load Classification Model
     try:
         print("    Downloading EfficientNet classification model...")
-        class_model_path = hf_hub_download(repo_id="sheikh987/efficientnet-isic-classifier", filename="efficientnet_isic_classifier_best.pth")
+        class_model_path = hf_hub_download(
+            repo_id="sheikh987/efficientnet-isic-classifier",
+            filename="efficientnet_isic_classifier_best.pth",
+            cache_dir="./models_cache"
+        )
         classification_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=7).to(DEVICE)
         classification_model.load_state_dict(torch.load(class_model_path, map_location=DEVICE))
         classification_model.eval()
@@ -70,11 +80,20 @@ def load_assets():
         print("    ✅ Knowledge base loaded.")
     except Exception as e:
         print(f"!!! FATAL: Could not load knowledge_base.json: {e}")
-        
+
     # Define Image Transforms
-    transform_segment = A.Compose([A.Resize(256, 256), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0), ToTensorV2()])
-    transform_classify = transforms.Compose([transforms.Resize((300, 300)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    
+    transform_segment = A.Compose([
+        A.Resize(256, 256),
+        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0),
+        ToTensorV2()
+    ])
+
+    transform_classify = transforms.Compose([
+        transforms.Resize((300, 300)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
     print("\n--> API is ready to accept requests.")
 
 # --- 2. DEFINE THE MAIN API ENDPOINT ---
@@ -101,7 +120,7 @@ async def analyze_image(request: ImageRequest):
     with torch.no_grad():
         seg_logits = segmentation_model(seg_input_tensor)
     seg_mask = (torch.sigmoid(seg_logits) > 0.5).float().squeeze().cpu().numpy()
-    
+
     if seg_mask.sum() < 200:
         return {"status": "Failed", "message": "No lesion could be clearly identified in the image."}
 
@@ -118,18 +137,17 @@ async def analyze_image(request: ImageRequest):
     with torch.no_grad():
         class_logits = classification_model(class_input_tensor)
         probabilities = torch.nn.functional.softmax(class_logits, dim=1)
-    
+
     confidence, predicted_idx = torch.max(probabilities, 1)
     confidence_percent = confidence.item() * 100
-    
-    # SAFETY NET
+
     if confidence_percent < 75.0:
-         return {"status": "Inconclusive", "message": f"Model confidence ({confidence_percent:.2f}%) is below the 75% threshold."}
+        return {"status": "Inconclusive", "message": f"Model confidence ({confidence_percent:.2f}%) is below the 75% threshold."}
 
     # STAGE 3: LOOKUP AND RETURN RESULT
     predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
     info = knowledge_base.get(predicted_abbr, {})
-    
+
     return {
         "status": "Success",
         "prediction": info,
@@ -140,4 +158,4 @@ async def analyze_image(request: ImageRequest):
 # --- Root endpoint to check if the API is alive ---
 @app.get("/")
 def root():
-    return {"message": "AI Skin Lesion Analyzer API is running."}
+    return {"message": "AI Skin Lesion Analyzer API is running."}