Update app.py

app.py

@@ -1,118 +1,82 @@
 import numpy as np
-import tempfile
 import pandas as pd
-import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, GlobalAveragePooling2D
-from tensorflow.keras.models import Model, load_model, Sequential
-from tensorflow.keras.optimizers import Adam
-from tensorflow.keras.metrics import Precision, Recall
-from tensorflow.keras.callbacks import EarlyStopping
-from sklearn.utils.class_weight import compute_class_weight
-from sklearn.model_selection import train_test_split
-from tensorflow.image import resize
 import cv2
-
-from tensorflow.keras.utils import to_categorical
-import matplotlib.pyplot as plt
-import warnings
-import warnings
-warnings.filterwarnings("ignore")
-
-# print ('modules loaded')
-
-
-import streamlit as st
-import pandas as pd 
-import numpy as np
+import torch
+from tensorflow import keras
 from PIL import Image
-import tensorflow.keras as keras
+from transformers import ViTForImageClassification, ViTFeatureExtractor
+from tensorflow.keras.models import load_model
+import streamlit as st
 
 st.title("Skin Cancer Classification App")
 
+# Load TensorFlow models
 models = {
-  "Le_Net": load_model('LeNet_5.h5'),
-  "Simple_CNN": load_model('Simple CNN.h5'),
-  "Alex_Net": load_model('AlexNet.h5'),
-  "Deeper_CNN": load_model('Deeper CNN.h5')
+    "Le_Net": load_model('LeNet_5.h5'),
+    "Simple_CNN": load_model('Simple CNN.h5'),
+    "Alex_Net": load_model('AlexNet.h5'),
+    "Deeper_CNN": load_model('Deeper CNN.h5'),
 }
 
+# Load PyTorch ViT model
+vit_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=7)
+vit_model.load_state_dict(torch.load('./vit_skin_cancer_model.pth'))
+vit_model.eval()  # Set the model to evaluation mode
+
+# Add the PyTorch model to the models dictionary
+models["ViT_Model"] = vit_model
+
 # Allow user to select model
 model_name = st.selectbox("Choose a model", list(models.keys()))
 model = models[model_name]
 
 # Upload Image
-
 file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
-# file ='hmnist_28_28_RGB.csv'
-
-print(file)
-
-true_file=pd.read_csv("HAM10000_metadata.csv")
-
-
-# true_file.apply(lambda x: x["image_id"] ==file)
 
-# uploaded_file = st.file_uploader("Choose a image file", type="jpg")
+true_file = pd.read_csv("HAM10000_metadata.csv")
 
-# if uploaded_file is not None:
-#     # Convert the file to an opencv image.
-#     file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-#     opencv_image = cv2.imdecode(file_bytes, 1)
-
-#     st.image(opencv_image, channels="BGR")
-    
-
-classes = {4: ('nv', ' melanocytic nevi'), 6: ('mel', 'melanoma'),
-           2 :('bkl', 'benign keratosis-like lesions'), 1:('bcc' , ' basal cell carcinoma'),
-           5: ('vasc', ' pyogenic granulomas and hemorrhage'),
-           0: ('akiec', 'Actinic keratoses and intraepithelial carcinomae'),
-           3: ('df', 'dermatofibroma')}
-
-classes_map = {'nv': 'melanocytic nevi', 
-               'mel': 'melanoma',
-           'bkl':'benign keratosis-like lesions',
-               'bcc':' basal cell carcinoma',
-           'vasc': ' pyogenic granulomas and hemorrhage',
-           'akiec': 'Actinic keratoses and intraepithelial carcinomae',
-           'df': 'dermatofibroma'}
+classes = {
+    4: ('nv', 'melanocytic nevi'), 
+    6: ('mel', 'melanoma'),
+    2: ('bkl', 'benign keratosis-like lesions'), 
+    1: ('bcc', 'basal cell carcinoma'),
+    5: ('vasc', 'pyogenic granulomas and hemorrhage'),
+    0: ('akiec', 'Actinic keratoses and intraepithelial carcinomae'),
+    3: ('df', 'dermatofibroma')
+}
 
+classes_map = {
+    'nv': 'melanocytic nevi', 
+    'mel': 'melanoma',
+    'bkl': 'benign keratosis-like lesions',
+    'bcc': 'basal cell carcinoma',
+    'vasc': 'pyogenic granulomas and hemorrhage',
+    'akiec': 'Actinic keratoses and intraepithelial carcinomae',
+    'df': 'dermatofibroma'
+}
 
 if file is not None:
-
     file_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
     opencv_image = cv2.imdecode(file_bytes, 1)
     
-    
-    # temp_dir = tempfile.TemporaryDirectory()
-    # temp_file_path = temp_dir.name + "/" + file.name
-    
-    # # Save the uploaded file to the temporary directory
-    # with open(temp_file_path, "wb") as f:
-    #     f.write(file.read())
-        
-    # img = cv2.imread(file)
-    # cv2_imshow(img)
+    # Resize image for TensorFlow models
     img1 = cv2.resize(opencv_image, (32, 32))
-    result = model.predict(img1.reshape(1, 32, 32, 3))
-    max_prob = max(result[0])
-    class_ind = list(result[0]).index(max_prob)
-    class_name = classes[class_ind]
-    # print(class_name)
-    # count+=1
-    # if count>10:
-    #     break
-
-        
-    # df = pd.read_csv(file)  
-    # # Get first row 
-    # img_reshaped = image_resize(df)
     
-    # # Get prediction
-    # pred = model.predict(img_reshaped)
-    # label = np.argmax(pred)
+    if model_name == "ViT_Model":
+        # PyTorch model inference
+        feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+        image = feature_extractor(images=opencv_image, return_tensors="pt")['pixel_values']
+        with torch.no_grad():
+            outputs = model(image)
+            class_ind = outputs.logits.argmax(-1).item()
+            class_name = classes[class_ind]
+    else:
+        # TensorFlow model inference
+        result = model.predict(img1.reshape(1, 32, 32, 3))
+        max_prob = max(result[0])
+        class_ind = list(result[0]).index(max_prob)
+        class_name = classes[class_ind]
 
-    
     # Display image and result 
     col1, col2 = st.columns(2)
     with col1:
@@ -125,38 +89,6 @@ if file is not None:
             name = file.name.split(".")[0]
         if name in true_file['image_id'].values: 
             st.write("True Label: ", classes_map[true_file.loc[true_file['image_id']==name, 'dx'].iloc[0]])
-            st.write("Prediction:",class_name[1])
+            st.write("Prediction:", class_name[1])
         else:
             st.write("No match")
-        
-        
-        # st.write(file.name)
-        # st.write("True Label",true_file[true_file.image_id==file.name]["dx"][0])
-        # st.write("Prediction",class_name[0])
-        # st.metric("Category:", class_name[1])
-
-
-        
-
-        
-
-
-
-
-
-# from google.colab.patches import cv2_imshow
-# srcdir = '/kaggle/input/skin-cancer-mnist-ham10000/HAM10000_images_part_1'
-# count=0
-# for temp in os.listdir(srcdir):
-#     img = cv2.imread(os.path.join(srcdir, temp))
-#     cv2.imwrite(temp, img)
-#     cv2_imshow(img)
-#     img = cv2.resize(img, (28, 28))
-#     result = model.predict(img.reshape(1, 28, 28, 3))
-#     max_prob = max(result[0])
-#     class_ind = list(result[0]).index(max_prob)
-#     class_name = classes[class_ind]
-#     print(class_name)
-#     count+=1
-#     if count>10:
-#         break