Update app.py

app.py

@@ -1,30 +1,34 @@
+import numpy as np
+from keras.models import load_model
+from keras.preprocessing import image
 import gradio as gr
-from transformers import AutoModelForImageClassification, AutoImageProcessor
-from PIL import Image
-import torch
 
-# Load the working public model
-model_name = "projectswithjam/skin-condition-classifier"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-processor = AutoImageProcessor.from_pretrained(model_name)
-labels = model.config.id2label
+# Load the model (ensure this .h5 file is in the same folder)
+model = load_model("VGG16-Final.h5")
+
+# Class labels (update if different)
+class_names = [
+    'Alopecia Areata', 'Contact Dermatitis', 'Folliculitis',
+    'Head Lice', 'Lichen Planus', 'Male Pattern Baldness',
+    'Psoriasis', 'Seborrheic Dermatitis', 'Telogen Effluvium',
+    'Tinea Capitis'
+]
 
 # Prediction function
-def predict(image):
-    image = image.convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_idx = logits.argmax(-1).item()
-    label = labels[str(predicted_idx)]
-    return f"Prediction: {label}"
+def predict(img):
+    img = img.resize((224, 224))  # Match input shape of model
+    img_array = image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0) / 255.0  # Normalize
+    prediction = model.predict(img_array)
+    predicted_class = class_names[np.argmax(prediction)]
+    return f"Prediction: {predicted_class}"
 
-# Gradio Interface
+# Gradio interface
 gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Skin Disease Classifier",
-    description="Upload a skin or scalp image. This model predicts the skin condition using a transformer trained on DermNet data."
+    title="Hair/Scalp Disease Classifier",
+    description="Upload a scalp image to classify the condition using a VGG16-based CNN model."
 ).launch()
+