Create main.py

main.py

@@ -0,0 +1,36 @@
+import streamlit as st
+from PIL import Image
+import torch
+
+from model import ModelColorization
+
+from utils import process_gs_image, inverse_transform_cs
+
+# create the model
+model = ModelColorization().from_pretrained("sebastiansarasti/AutoEncoderImageColorization")
+
+# create the streamlit app
+st.title("Image Colorization App")
+st.write("This is an app to colorize black and white images.")
+
+# create a botton to upload the image
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+# check if the image is uploaded
+if uploaded_file is not None:
+    # display the image
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image.", use_container_width=True)
+
+    # create a button to colorize the image
+    if st.button("Colorize"):
+        # process the grayscale image
+        image, original_size = process_gs_image(image)
+        # run the model
+        model.eval()
+        with torch.no_grad():
+            result = model(image)
+        # colorize the image
+        colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
+        # display the colorized image
+        st.image(colorized_image, caption="Colorized Image.", use_container_width=True)