Update main.py

main.py

@@ -7,16 +7,88 @@ os.environ["XDG_CACHE_HOME"] = os.path.join(os.getcwd(), ".cache")
 from fastapi import FastAPI, File, UploadFile,Form
 from fastapi.responses import FileResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
+import PIL
+import cv2
+import numpy as np
+import uuid
+from zipfile import ZipFile, ZIP_DEFLATED
+from io import BytesIO
+from random import randint
+from datetime import datetime
+
 from src.deoldify import device
 from src.deoldify.device_id import DeviceId
+from src.deoldify.visualize import *
 from src.app_utils import get_model_bin
-from colorize_image_function import colorize_image
 
-app = FastAPI()
 
 device.set(device=DeviceId.CPU)
-model_dir = 'models/'
-colorizer = load_model(model_dir, "Artistic")
+
+
+@st.cache_resource
+def load_model(model_dir, option):
+    if option.lower() == 'artistic':
+        model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
+        get_model_bin(model_url, os.path.join(model_dir, "ColorizeArtistic_gen.pth"))
+        colorizer = get_image_colorizer(artistic=True)
+    elif option.lower() == 'stable':
+        model_url = "https://www.dropbox.com/s/usf7uifrctqw9rl/ColorizeStable_gen.pth?dl=0"
+        get_model_bin(model_url, os.path.join(model_dir, "ColorizeStable_gen.pth"))
+        colorizer = get_image_colorizer(artistic=False)
+
+    return colorizer
+
+
+def resize_img(input_img, max_size):
+    img = input_img.copy()
+    img_height, img_width = img.shape[0],img.shape[1]
+
+    if max(img_height, img_width) > max_size:
+        if img_height > img_width:
+            new_width = img_width*(max_size/img_height)
+            new_height = max_size
+            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
+            return resized_img
+
+        elif img_height <= img_width:
+            new_width = img_height*(max_size/img_width)
+            new_height = max_size
+            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
+            return resized_img
+
+    return img
+
+
+def colorize_image(pil_image, img_size=800) -> "PIL.Image":
+    # Open the image
+    pil_img = pil_image.convert("RGB")
+    img_rgb = np.array(pil_img)
+    resized_img_rgb = resize_img(img_rgb, img_size)
+    resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
+
+    # Send the image to the model
+    output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
+    
+    return output_pil_img
+
+
+def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
+    if fmt not in ["jpg", "png"]:
+        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
+    
+    pil_format = "JPEG" if fmt == "jpg" else "PNG"
+    file_format = "jpg" if fmt == "jpg" else "png"
+    mime = "image/jpeg" if fmt == "jpg" else "image/png"
+    
+    buf = BytesIO()
+    pil_image.save(buf, format=pil_format)
+    
+    return st.download_button(
+        label=label,
+        data=buf.getvalue(),
+        file_name=f'{filename}.{file_format}',
+        mime=mime,
+    )
 
 @app.post("/upload/")
 async def upload_file(file: UploadFile = File(...)):