Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
 import torch
 
-# Load the Hugging Face model and feature extractor
+# Cache the model and feature extractor
 @st.cache_resource
 def load_model():
     model_name = "syaha/skin_cancer_detection_model"
@@ -13,54 +13,55 @@ def load_model():
 
 model, feature_extractor = load_model()
 
-# App title
-st.title("Skin Cancer Detection Application")
+# App title and instructions
+st.title("Skin Cancer Detection App")
 st.write("""
-Upload an image of the affected skin area, and this app will identify and classify it based on pre-trained skin lesion types.
+Upload an image of the affected skin area, and this app will classify it based on pre-trained skin lesion types.
 """)
 
-# File uploader
+# File uploader for user input
 uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    # Display the uploaded image
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Perform inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        predicted_class = probabilities.argmax().item()
-        confidence = probabilities.max().item()
+    try:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        # Preprocess the image
+        inputs = feature_extractor(images=image, return_tensors="pt")
+        
+        # Perform inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+            predicted_class = probabilities.argmax().item()
+            confidence = probabilities.max().item()
 
-    # Map labels (update based on the model's labels)
-    labels = {
-        0: "Melanocytic nevi",
-        1: "Melanoma",
-        2: "Benign keratosis-like lesions",
-        3: "Basal cell carcinoma",
-        4: "Actinic keratoses",
-        5: "Vascular lesions",
-        6: "Dermatofibroma"
-    }
-    diagnosis = labels.get(predicted_class, "Unknown")
+        # Map predicted class to human-readable labels
+        labels = {
+            0: "Melanocytic nevi",
+            1: "Melanoma",
+            2: "Benign keratosis-like lesions",
+            3: "Basal cell carcinoma",
+            4: "Actinic keratoses",
+            5: "Vascular lesions",
+            6: "Dermatofibroma"
+        }
+        diagnosis = labels.get(predicted_class, "Unknown")
 
-    # Display results
-    st.write("### Results")
-    st.write(f"**Predicted Class**: {diagnosis}")
-    st.write(f"**Confidence**: {confidence:.2f}")
-    
-    # Generate report
-    st.write("### Report")
-    st.write(f"**Diagnosis**: {diagnosis}")
-    st.write(f"**Confidence Level**: {confidence:.2%}")
-    st.download_button(
-        label="Download Report",
-        data=f"Diagnosis: {diagnosis}\nConfidence Level: {confidence:.2%}",
-        file_name="report.txt",
-        mime="text/plain",
-    )
+        # Display results
+        st.subheader("Results")
+        st.write(f"**Diagnosis:** {diagnosis}")
+        st.write(f"**Confidence Level:** {confidence:.2%}")
+
+        # Provide an option to download the report
+        report = f"Diagnosis: {diagnosis}\nConfidence Level: {confidence:.2%}"
+        st.download_button(
+            label="Download Report",
+            data=report,
+            file_name="report.txt",
+            mime="text/plain",
+        )
+    except Exception as e:
+        st.error(f"An error occurred while processing the image: {e}")