Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, File, UploadFile, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import os | |
| import uuid | |
| import tempfile | |
| app = FastAPI( | |
| title="Medical Image Classification API", | |
| description="AI-powered medical image classification service", | |
| version="1.0.0" | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| model = None | |
| class_names = { | |
| 0: "AKIEC", # Actinic keratoses and intraepithelial carcinoma | |
| 1: "BCC", # Basal cell carcinoma | |
| 2: "BKL", # Benign keratosis-like lesions | |
| 3: "DF", # Dermatofibroma | |
| 4: "MEL", # Melanoma | |
| 5: "NV", # Melanocytic nevi | |
| 6: "VASC" # Vascular lesions | |
| } | |
| full_names = { | |
| "AKIEC": "Actinic keratoses and intraepithelial carcinoma", | |
| "BCC": "Basal cell carcinoma", | |
| "BKL": "Benign keratosis-like lesions", | |
| "DF": "Dermatofibroma", | |
| "MEL": "Melanoma", | |
| "NV": "Melanocytic nevi", | |
| "VASC": "Vascular lesions" | |
| } | |
| UPLOAD_DIR = tempfile.mkdtemp() | |
| def load_model(): | |
| global model | |
| model_path = "efficientnetv2s.h5" | |
| if not os.path.exists(model_path): | |
| raise FileNotFoundError(f"Model file not found: {model_path}") | |
| model = tf.keras.models.load_model(model_path) | |
| return model | |
| def predict_image(image_path): | |
| global model | |
| if model is None: | |
| model = load_model() | |
| img = Image.open(image_path).convert('RGB') | |
| img = img.resize((224, 224)) | |
| img_array = np.array(img) | |
| img_array = np.expand_dims(img_array, axis=0) | |
| img_array = img_array.astype(np.float32) / 255.0 | |
| predictions = model.predict(img_array) | |
| predicted_class = np.argmax(predictions[0]) | |
| confidence = float(predictions[0][predicted_class]) | |
| class_code = class_names[predicted_class] | |
| class_full_name = full_names[class_code] | |
| all_predictions = [] | |
| for i, prob in enumerate(predictions[0]): | |
| all_predictions.append({ | |
| "label": class_names[i], | |
| "confidence": float(prob) | |
| }) | |
| all_predictions.sort(key=lambda x: x["confidence"], reverse=True) | |
| return class_full_name, confidence, all_predictions | |
| async def health_check(): | |
| return {"status": "healthy", "service": "medical-image-classifier"} | |
| async def classify_image_api(file: UploadFile = File(...)): | |
| try: | |
| if not file.content_type or not file.content_type.startswith('image/'): | |
| raise HTTPException(status_code=400, detail="File must be an image") | |
| unique_filename = f"{uuid.uuid4().hex}_{file.filename}" | |
| file_path = os.path.join(UPLOAD_DIR, unique_filename) | |
| with open(file_path, "wb") as buffer: | |
| content = await file.read() | |
| buffer.write(content) | |
| label, confidence, all_predictions = predict_image(file_path) | |
| os.remove(file_path) | |
| formatted_predictions = [] | |
| for pred in all_predictions: | |
| formatted_predictions.append({ | |
| "label": pred["label"], | |
| "confidence": float(pred["confidence"]), | |
| "confidence_percent": f"{pred['confidence'] * 100:.2f}%" | |
| }) | |
| return JSONResponse( | |
| status_code=200, | |
| content={ | |
| "success": True, | |
| "prediction": { | |
| "top_prediction": { | |
| "label": label, | |
| "confidence": float(confidence), | |
| "confidence_percent": f"{confidence * 100:.2f}%" | |
| }, | |
| "all_predictions": formatted_predictions | |
| } | |
| } | |
| ) | |
| except Exception as e: | |
| if 'file_path' in locals() and os.path.exists(file_path): | |
| os.remove(file_path) | |
| raise HTTPException(status_code=500, detail=f"Classification failed: {str(e)}") | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=8003) | |