Update app.py

app.py

@@ -39,44 +39,89 @@ 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()
+left_column, right_column = st.columns(2)
 
+with left_column:
 # Title of the app
-st.title("Image Upload and Inference Example")
-
-# File uploader for images
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    # Open and display the image
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image.', use_column_width=True)
-    
-    # Perform inference
-    st.write("Performing inference...")
+    st.title("Original Model")
+    # File uploader for images
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
     
-    # 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)
+    if uploaded_file is not None:
+        # Open and display the image
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image.', width=300)
+        
+        # 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}%")
+def extract(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)
+        
+with right_column:
+    uploaded_file = st.file_uploader("Upload ZIP with images for celebrity to forget.", type="zip")
+        if uploaded_file is not None:
+        st.write("Uploaded ZIP file details:")
+        st.write({
+            "Filename": uploaded_file.name,
+        })
 
-    freq = Counter(preds)
-    top_three = freq.most_common(3)
-    for celeb, count in top_three:
-        st.write(f"{celeb}: {int(count)*2}%")
+        # Call function to extract and display images
+        extract(uploaded_file)
+        st.write("Unlearning begins...")
+        unlearn()
+        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])