IvaElen commited on
Commit
8781c55
1 Parent(s): b9e8251

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +16 -11
main.py CHANGED
@@ -1,24 +1,29 @@
1
- import streamlit as st
 
2
 
3
  import pandas as pd
 
 
4
 
5
  import clip
6
- import torchvision.transforms as transforms
7
- import torchvision.datasets as datasets
8
  import torch
9
- import numpy as np
10
- import random
11
- from get_similiarty import get_similiarity
12
 
13
- device = "cuda" if torch.cuda.is_available() else "cpu"
14
- #load model -resnet50
15
 
16
 
17
- model_resnet = torch.load("model.pt", device )
18
 
 
 
19
  #load model - ViT-B/32
20
- model_vit = torch.load("model_vit.pt", device )
21
 
 
 
 
 
 
22
 
23
  st.title('Find my pic!')
24
 
@@ -53,4 +58,4 @@ if txt is not None:
53
  st.write(img_desc[ind])
54
  with col2:
55
  st.image(pic[1])
56
- st.write(img_desc[ind])
 
1
+ import zipfile
2
+ import random
3
 
4
  import pandas as pd
5
+ import numpy as np
6
+ import streamlit as st
7
 
8
  import clip
 
 
9
  import torch
10
+ import torchvision.transforms as transforms
 
 
11
 
12
+ from get_similiarty import get_similiarity
 
13
 
14
 
15
+ device = "cuda" if torch.cuda.is_available() else "cpu"
16
 
17
+ #load model -resnet50
18
+ model_resnet = torch.load("model.pt", device)
19
  #load model - ViT-B/32
20
+ model_vit = torch.load("model_vit.pt", device)
21
 
22
+ #Распаковка ZIP-файла с фотографиями
23
+ zip_file_path = "sample.zip"
24
+ target_folder = "sample/"
25
+ with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
26
+ zip_ref.extractall(target_folder)
27
 
28
  st.title('Find my pic!')
29
 
 
58
  st.write(img_desc[ind])
59
  with col2:
60
  st.image(pic[1])
61
+ st.write(img_desc[ind])