Update app.py

app.py

@@ -9,6 +9,7 @@ import timm
 from tqdm import tqdm
 import torch.nn.functional as F 
 from collections import Counter
+from scrub import unlearn
 # Your model and other necessary functions
 # Assuming model, device, final_conv, fc_params, cls_names, and any other needed components are defined elsewhere
 # from your_model_module import model, device, final_conv, fc_params, cls_names, SaveFeatures, getCAM, tensor_2_im
@@ -39,27 +40,11 @@ reversed_map = {
     16: 'Will Smith'
 }
 
-def extract_and_display_images(zip_file):
-    # Create a directory to store extracted images
-    extract_path = "extracted_images"
-    os.makedirs(extract_path, exist_ok=True)
-
-    # Extract images from the ZIP file
-    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
-        zip_ref.extractall(extract_path)
-
-    # Display each image in the extracted directory
-    image_files = os.listdir(extract_path)
-    for image_file in image_files:
-        image_path = os.path.join(extract_path, image_file)
-        image = Image.open(image_path)
-        st.image(image, caption=image_file, use_column_width=True)
-        
 model = timm.create_model("rexnet_150", pretrained = True, num_classes = 17)  
 model.load_state_dict(torch.load('faces_best_model.pth', map_location=torch.device('cpu')))
 model.eval()
+extract('celeb-dataset.zip', 'celeb-dataset')
 left_column, right_column = st.columns(2)
-
 with left_column:
 # Title of the app
     st.title("Original Model")
@@ -97,9 +82,8 @@ with left_column:
         top_three = freq.most_common(3)
         for celeb, count in top_three:
             st.write(f"{celeb}: {int(count)*2}%")
-def extract(zip_file):
+def extract(zip_file, extract_path):
     # Create a directory to store extracted images
-    extract_path = "extracted_images"
     os.makedirs(extract_path, exist_ok=True)
 
     # Extract images from the ZIP file
@@ -114,14 +98,41 @@ def extract(zip_file):
         
 with right_column:
     uploaded_file = st.file_uploader("Upload ZIP with images for celebrity to forget.", type="zip")
-        if uploaded_file is not None:
+    if uploaded_file is not None:
         st.write("Uploaded ZIP file details:")
         st.write({
             "Filename": uploaded_file.name,
         })
 
         # Call function to extract and display images
-        extract(uploaded_file)
+        extract(uploaded_file, 'forget_set')
         st.write("Unlearning begins...")
         unlearn()
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+                # Perform inference
+        st.write("Performing inference...")
+        
+        # Transform the image to fit model requirements
+        preprocess = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        
+        image_tensor = preprocess(image).unsqueeze(0)
+        
+        preds = []
+        with torch.no_grad():
+            for i in range(50): 
+                output = model(image_tensor)
+                probabilities = F.softmax(output, dim=1)
+                pred_class = torch.argmax(probabilities, dim=1)
+                pred_label = reversed_map[pred_class.item()]
+                preds.append(pred_label)
+        
+    
+        freq = Counter(preds)
+        top_three = freq.most_common(3)
+        for celeb, count in top_three:
+            st.write(f"{celeb}: {int(count)*2}%")
+