update output size

app.py

@@ -52,12 +52,16 @@ def preprocess_image(image):
     image = transforms.ToTensor()(image)[:1] * 2. - 1.
     return image
 
-def postprocess_image(grayscale, prediction):
-    return lab_to_rgb(grayscale.unsqueeze(0), prediction.cpu())[0]
+def postprocess_image(grayscale, prediction, original_size):
+    # Convert Lab back to RGB and resize to the original image size
+    colorized_image = lab_to_rgb(grayscale.unsqueeze(0), prediction.cpu())[0]
+    colorized_image = Image.fromarray((colorized_image * 255).astype("uint8"))
+    return colorized_image.resize(original_size)
 
 # Prediction function with output control
 def colorize_image(input_image, mode):
     grayscale_image = Image.fromarray(input_image).convert('L')
+    original_size = grayscale_image.size  # Store original size
     grayscale = preprocess_image(grayscale_image).to(device)
     
     with torch.no_grad():
@@ -65,9 +69,10 @@ def colorize_image(input_image, mode):
         mobilenet_output = mobilenet_model.net_G(grayscale.unsqueeze(0))
         autoencoder_output = autoencoder_model(grayscale.unsqueeze(0))
     
-    resnet_colorized = postprocess_image(grayscale, resnet_output)
-    mobilenet_colorized = postprocess_image(grayscale, mobilenet_output)
-    autoencoder_colorized = postprocess_image(grayscale, autoencoder_output)
+    # Resize outputs to match the original size
+    resnet_colorized = postprocess_image(grayscale, resnet_output, original_size)
+    mobilenet_colorized = postprocess_image(grayscale, mobilenet_output, original_size)
+    autoencoder_colorized = postprocess_image(grayscale, autoencoder_output, original_size)
     
     if mode == "ResNet":
         return resnet_colorized, None, None
@@ -117,7 +122,6 @@ def gradio_interface():
         )
 
         # Submit logic
-
         submit_button.click(
             fn=colorize_image,
             inputs=[input_image, output_modes],
@@ -127,7 +131,6 @@ def gradio_interface():
     return demo
 
 
-
 # Launch
 if __name__ == "__main__":
-    gradio_interface().launch()
+    gradio_interface().launch()

weight/denoising diffusion.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36234b7f753bfaa3c957c2f9cd67e0f4408eb29e479afb984c4df5fd6459d147
+size 124234454