EgorDan sunny-annie commited on
Commit
75fa944
1 Parent(s): 4f4712b

Upload 5 files (#1)

Browse files

- Upload 5 files (f288e14dd276e61e66abf345716434c0b84118e3)


Co-authored-by: Anya Baydenko <[email protected]>

pages/feature_extractor_v3.joblib ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6de31b46e55b824d28a5daab95de36f12f7cee1600bda97c7496433415c425c0
3
+ size 361
pages/img-2-txt.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import torch
3
+ from joblib import load
4
+ from PIL import Image
5
+
6
+ device = 'cpu'
7
+
8
+ tokenizer = load("./pages/tokenizer_v3.joblib")
9
+ feature_extractor = load("./pages/feature_extractor_v3.joblib")
10
+ # tokenizer = load("tokenizer_v3.joblib")
11
+ # feature_extractor = load("feature_extractor_v3.joblib")
12
+
13
+ model = load("./pages/img2txt_v4.joblib")
14
+ model.load_state_dict(torch.load("./pages/model_weights_i2t_fin.pt", map_location=torch.device('cpu')))
15
+
16
+ max_length = 512
17
+ min_length = 32
18
+ num_beams = 7
19
+ gen_kwargs = {"max_length": max_length, "min_length": min_length, "num_beams": num_beams}
20
+
21
+ uploaded_file = st.file_uploader("Выберите изображение обложки книги в формате jpeg или jpg...", type=["jpg", "jpeg"])
22
+
23
+ if uploaded_file is not None:
24
+ button = st.button('Сгенерировать описание')
25
+ image = Image.open(uploaded_file)
26
+ st.image(image, caption='Загруженное изображение')
27
+ image = image.resize([224,224])
28
+ if image.mode != "RGB":
29
+ image = image.convert(mode="RGB")
30
+
31
+
32
+ if button:
33
+ pixel_values = feature_extractor(images=[image], return_tensors="pt").pixel_values
34
+ pixel_values = pixel_values.to(device)
35
+ output_ids = model.generate(pixel_values, **gen_kwargs)
36
+ preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
37
+ preds = [pred.strip() for pred in preds]
38
+ st.write(preds[0])
pages/img2txt_v4.joblib ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cb9bc00567a3ae345d9877bc7139e6311f45bd1360f58bd5f027967da67d3789
3
+ size 1171173081
pages/model_weights_i2t_fin.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6753b600e5d753ad7420a92aa4dcd3f5d860f8e9b5c933f059144e4e579d938e
3
+ size 1171154606
pages/tokenizer_v3.joblib ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2cc0b03a6f6232c82c5d9afa89f8d53f71d1b324e9b3845869d4871fa0ebe87d
3
+ size 2617590