Spaces:
Sleeping
Sleeping
File size: 5,412 Bytes
1994e69 1ff265e 48673e5 d165ddc 7c12725 46489a3 7c12725 46489a3 d1b1c42 38412b8 763c5ad bb5701d 7c12725 496c703 1657dc5 496c703 1372c37 1ac03ab 7c12725 38412b8 8a811c4 7c12725 8a811c4 7c12725 8a811c4 7c12725 8a811c4 abd119a 8a811c4 38412b8 d1b1c42 8a811c4 38412b8 e40a05a bb5701d 0a3c48d 2945403 38412b8 0a3c48d 38412b8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import sys
import os
sys.path.append(os.path.dirname(__file__))
import streamlit as st
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2
from torchvision import transforms
import timm
from tqdm import tqdm
import torch.nn.functional as F
from collections import Counter
from scrub import unlearn
import zipfile
import time
# 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
# Set up your model and device here
# model = ...
# device = ...
# final_conv = ...
# fc_params = ...
# cls_names = ...
reversed_map = {
0: 'Angelina Jolie',
1: 'Brad Pitt',
2: 'Denzel Washington',
3: 'Hugh Jackman',
4: 'Jennifer Lawrence',
5: 'Johnny Depp',
6: 'Kate Winslet',
7: 'Leonardo DiCaprio',
8: 'Megan Fox',
9: 'Natalie Portman',
10: 'Nicole Kidman',
11: 'Robert Downey Jr',
12: 'Sandra Bullock',
13: 'Scarlett Johansson',
14: 'Tom Cruise',
15: 'Tom Hanks',
16: 'Will Smith'
}
def extract(zip_file, extract_path):
os.makedirs(extract_path, exist_ok=True)
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
zip_ref.extractall(extract_path)
for root, dirs, files in os.walk(extract_path):
for file in files:
if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')): # Adjust file extensions as needed
image_path = os.path.join(root, file)
try:
image = Image.open(image_path)
# Process or display the image here
except IOError as e:
print(f"Error opening image {image_path}: {e}")
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")
# 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.', 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}%")
with right_column:
st.title("Unlearned Model")
options = ['SSD', 'SCRUB', 'UNSIR', 'Incompetent Teacher', 'Mislabel']
# Display dropdown and store selected value
selected_option = st.selectbox('Select an option:', options)
uploaded_file = st.file_uploader("Upload ZIP with images for celebrity to forget.", type="zip")
if uploaded_file is not None:
# Call function to extract and display images
extract(uploaded_file, 'forget_set')
st.write("Unlearning...")
#unlearn()
time.sleep(10)
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)
# 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_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}%")
|