Update app.py

app.py

@@ -116,39 +116,41 @@ with right_column:
         extract(uploaded_file, 'forget_set')
         st.write("Unlearning...")
         #unlearn()
-        time.sleep(10)
+        time.sleep(5)
         model_s = timm.create_model("rexnet_150", pretrained = True, num_classes = 17)  
         model_s.load_state_dict(torch.load('celeb-model-unlearned.pth', map_location=torch.device('cpu')))
         model_s.eval()
 
         uploaded_file2 = st.file_uploader("Choose image...", type=["jpg", "jpeg", "png"])
-        image2 = Image.open(uploaded_file2)
-        st.image(image2, caption='Uploaded Image.', width=300)        
+        if uploaded_file2 is not None:
+
+            image2 = Image.open(uploaded_file2)
+            st.image(image2, 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]),
-        ])
+            st.write("Performing inference...")
         
-        image_tensor = preprocess(image).unsqueeze(0)
+            # 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(image2).unsqueeze(0)
+            
+            preds = []
+            with torch.no_grad():
+                for i in range(50): 
+                    output = model_s(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)
+            
         
-        preds = []
-        with torch.no_grad():
-            for i in range(50): 
-                output = model_s(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}%")
+            freq = Counter(preds)
+            top_three = freq.most_common(3)
+            for celeb, count in top_three:
+                st.write(f"{celeb}: {int(count)*2}%")