Update main.py

main.py

@@ -1,48 +1,53 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-import tensorflow as tf
-import numpy as np
-from vit_keras import vit
-import tensorflow_addons as tfa
-from io import BytesIO
-from PIL import Image
-
-app = FastAPI()
-
-# CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Model setup
-vit_model = vit.vit_b16(image_size=224, activation='softmax', pretrained=True, include_top=False, pretrained_top=False, classes=7)
-model1 = tf.keras.Sequential([
-    vit_model,
-    tf.keras.layers.Flatten(),
-    tf.keras.layers.BatchNormalization(),
-    tf.keras.layers.Dense(11, activation=tfa.activations.gelu),
-    tf.keras.layers.BatchNormalization(),
-    tf.keras.layers.Dense(7, activation='softmax')
-])
-model1.load_weights('vit_model_weights.h5')
-
-labels = ['akiec', 'bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
-
-def preprocess_image(image: Image.Image):
-    image = image.resize((224, 224))
-    image = np.array(image) / 255.0
-    image = np.expand_dims(image, axis=0)
-    return image
-
-@app.post("/predict/")
-async def predict(file: UploadFile = File(...)):
-    image = Image.open(BytesIO(await file.read()))
-    processed_image = preprocess_image(image)
-    predictions = model1.predict(processed_image)
-    predicted_class = labels[np.argmax(predictions)]
-    confidence = np.max(predictions)
-    return {"predicted_class": predicted_class, "confidence": float(confidence)}
+# main.py
+from fastapi import FastAPI, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+import tensorflow as tf
+import numpy as np
+from vit_keras import vit
+import tensorflow_addons as tfa
+from io import BytesIO
+from PIL import Image
+
+app = FastAPI()
+
+# CORS configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust this according to your security needs
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load the model
+vit_model = vit.vit_b16(image_size=224, activation='softmax', pretrained=True, include_top=False, pretrained_top=False, classes=7)
+model1 = tf.keras.Sequential([
+    vit_model,
+    tf.keras.layers.Flatten(),
+    tf.keras.layers.BatchNormalization(),
+    tf.keras.layers.Dense(11, activation=tfa.activations.gelu),
+    tf.keras.layers.BatchNormalization(),
+    tf.keras.layers.Dense(7, activation='softmax')
+])
+model1.load_weights('vit_model_weights.h5')
+
+labels = ['akiec', 'bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
+
+def preprocess_image(image: Image.Image):
+    image = image.resize((224, 224))
+    image = np.array(image) / 255.0
+    image = np.expand_dims(image, axis=0)
+    return image
+
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    image = Image.open(BytesIO(await file.read()))
+    processed_image = preprocess_image(image)
+    predictions = model1.predict(processed_image)
+    predicted_class = labels[np.argmax(predictions)]
+    confidence = np.max(predictions)
+    return {"predicted_class": predicted_class, "confidence": float(confidence)}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)