barunsaha commited on
Commit
37905b9
1 Parent(s): 4cb9b95

Fix #1: only use uploaded image files

Browse files
Files changed (1) hide show
  1. app.py +9 -34
app.py CHANGED
@@ -1,6 +1,5 @@
1
  import logging
2
  import os
3
- import random
4
 
5
  import PIL
6
  import requests
@@ -55,17 +54,11 @@ SAFETY_SETTINGS = [
55
  }
56
  ]
57
 
58
- SAMPLE_IMAGE_URLS = [
59
- 'https://upload.wikimedia.org/wikipedia/commons/thumb/1/10/Websocket_connection.png/220px-Websocket_connection.png',
60
- 'https://assets-global.website-files.com/5ff66329429d880392f6cba2/63fe48adb8834a29a618ce84_148.3.png',
61
- 'https://media.geeksforgeeks.org/wp-content/uploads/bus1.png',
62
- ]
63
-
64
 
65
  @st.cache_resource
66
  def get_gemini_model():
67
  """
68
- Get the Gemini Pro Vision model.
69
 
70
  :return: The model
71
  """
@@ -98,7 +91,7 @@ genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
98
 
99
  st.title('Sys2Doc')
100
  st.header(
101
- 'Generate Documentation Based on System Diagram—powered by Gemini 1.5 Flash'
102
  )
103
 
104
  uploaded_file = st.file_uploader(
@@ -107,36 +100,18 @@ uploaded_file = st.file_uploader(
107
  type=SUPPORTED_FILE_EXTENSIONS
108
  )
109
 
110
- st.write('OR provide the URL of the image:')
111
- img_url = st.text_input(
112
- label='URL of the image',
113
- value=random.choice(SAMPLE_IMAGE_URLS),
114
- )
115
- st.markdown(
116
- '(*If an image is uploaded and a URL is also provided,'
117
- ' Sys2Doc will consider the uploaded image*)'
118
- )
119
-
120
- if uploaded_file is not None or (img_url is not None and len(img_url) > 0):
121
  # Show the uploaded image & related info
122
  the_img = None
123
  file_details = None
124
 
125
  try:
126
- if uploaded_file:
127
- the_img = PIL.Image.open(uploaded_file)
128
- file_details = {
129
- 'file_name': uploaded_file.name,
130
- 'file_type': uploaded_file.type,
131
- 'file_size': uploaded_file.size
132
- }
133
- elif img_url:
134
- the_img = PIL.Image.open(requests.get(img_url, stream=True).raw)
135
- file_details = {
136
- 'file_name': os.path.basename(img_url),
137
- 'file_type': the_img.format,
138
- # 'file_info': the_img.info
139
- }
140
 
141
  if the_img and the_img.mode in ('RGBA', 'P'):
142
  the_img = the_img.convert('RGB')
 
1
  import logging
2
  import os
 
3
 
4
  import PIL
5
  import requests
 
54
  }
55
  ]
56
 
 
 
 
 
 
 
57
 
58
  @st.cache_resource
59
  def get_gemini_model():
60
  """
61
+ Get the Gemini model.
62
 
63
  :return: The model
64
  """
 
91
 
92
  st.title('Sys2Doc')
93
  st.header(
94
+ 'Generate Documentation Based on System Diagram powered by Gemini 1.5 Flash'
95
  )
96
 
97
  uploaded_file = st.file_uploader(
 
100
  type=SUPPORTED_FILE_EXTENSIONS
101
  )
102
 
103
+ if uploaded_file is not None:
 
 
 
 
 
 
 
 
 
 
104
  # Show the uploaded image & related info
105
  the_img = None
106
  file_details = None
107
 
108
  try:
109
+ the_img = PIL.Image.open(uploaded_file)
110
+ file_details = {
111
+ 'file_name': uploaded_file.name,
112
+ 'file_type': uploaded_file.type,
113
+ 'file_size': uploaded_file.size
114
+ }
 
 
 
 
 
 
 
 
115
 
116
  if the_img and the_img.mode in ('RGBA', 'P'):
117
  the_img = the_img.convert('RGB')