Update pages/4_ImageNet1k.py

pages/4_ImageNet1k.py

@@ -32,17 +32,23 @@ with col2:
     image = None
     with st.form("display image"):
         img_index = st.text_input('Image ID to display')
-        try:
-            img_index = int(img_index)
-        except:
-            pass
-
+    
         submitted = st.form_submit_button('Display this image')
+        error_container = st.empty()
+
         if submitted:
-            image = dataset_dict[img_index//10_000][img_index%10_000]['image']
-            class_label = dataset_dict[img_index//10_000][img_index%10_000]['label']
-            class_id = dataset_dict[img_index//10_000][img_index%10_000]['id']
+            try:
+                img_index = int(img_index)
+                if img_index > 50000 or img_index < 0:
+                    error_container.error('The Image ID must be in range from 0 to 50000', icon="🚫")
+                else:
+                    image = dataset_dict[img_index//10_000][img_index%10_000]['image']
+                    class_label = dataset_dict[img_index//10_000][img_index%10_000]['label']
+                    class_id = dataset_dict[img_index//10_000][img_index%10_000]['id']
+            except ValueError:
+                error_container.error('Please enter an integer number for Image ID', icon = "🚫")
+            
     if image != None:
         st.image(image)
         st.write('**Class label:** ', class_label)
-        st.write('\n**Class id:** ', str(class_id))
+        st.write('\n**Class id:** ', str(class_id))