5

app.py

@@ -134,8 +134,8 @@ def main():
 
             for uploaded_file in uploaded_files:
                 try:
-                    st.image(uploaded_file, caption=f"Uploaded Image: {uploaded_file.name}", use_container_width=True)
                     file_path = save_uploaded_image(uploaded_file)
+                    st.image(file_path, caption=f"Uploaded Image: {uploaded_file.name}", use_column_width=True)
 
                     prediction = classify_image(file_path)
                     results[file_path] = prediction
@@ -160,8 +160,8 @@ def main():
 
             for training_file in training_files:
                 try:
-                    st.image(training_file, caption=f"Uploaded Training Image: {training_file.name}", use_container_width=True)
                     file_path = save_uploaded_image(training_file)
+                    st.image(file_path, caption=f"Uploaded Training Image: {training_file.name}", use_column_width=True)
 
                     label = st.radio(
                         f"Classify {training_file.name}", ["Human", "Alien"], index=0, key=training_file.name
@@ -203,10 +203,10 @@ def main():
             if human_images:
                 st.info("These images have already been trained.")
                 for image_path in human_images:
-                    try:
-                        st.image(image_path, caption=f"Image: {os.path.basename(image_path)}", use_container_width=True)
-                    except Exception as e:
-                        st.error(f"Error displaying image {os.path.basename(image_path)}: {e}")
+                    if os.path.exists(image_path):
+                        st.image(image_path, caption=f"Image: {os.path.basename(image_path)}", use_column_width=True)
+                    else:
+                        st.error(f"File not found: {os.path.basename(image_path)}")
             else:
                 st.warning("No human images found for training.")
 
@@ -219,10 +219,10 @@ def main():
             if alien_images:
                 st.info("These images have already been trained.")
                 for image_path in alien_images:
-                    try:
-                        st.image(image_path, caption=f"Image: {os.path.basename(image_path)}", use_container_width=True)
-                    except Exception as e:
-                        st.error(f"Error displaying image {os.path.basename(image_path)}: {e}")
+                    if os.path.exists(image_path):
+                        st.image(image_path, caption=f"Image: {os.path.basename(image_path)}", use_column_width=True)
+                    else:
+                        st.error(f"File not found: {os.path.basename(image_path)}")
             else:
                 st.warning("No alien images found for training.")