google-research / maxim

[CVPR 2022 Oral] Official repository for "MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.
https://arxiv.org/abs/2201.02973
Apache License 2.0
977 stars 105 forks source link

Colab error when changing pre-trained model MAXIM-3S #43

Open MoraAlex opened 11 months ago

MoraAlex commented 11 months ago

In this cell https://colab.research.google.com/drive/1SjIBkPlBJrGFdIBqx7Lggcb-1H86Vq1h#scrollTo=Fcp68HNFf2Fy&line=3&uniqifier=1 when I change weight_drive_path to a model trained by MAXIM-3S I get the error in this cell https://colab.research.google.com/drive/1SjIBkPlBJrGFdIBqx7Lggcb-1H86Vq1h#scrollTo=Q1T544sZan4d&line=20&uniqifier=1 This is the error:

AttributeError                            Traceback (most recent call last)

[<ipython-input-14-ede7aca02ff0>](https://localhost:8080/#) in <cell line: 6>()
      5 
      6 for path in images:
----> 7   im = Image.fromarray(predict(path))
      8   im.save("/content/gdrive/MyDrive/maxim/output/"+path.split('/')[-1])
      9 

1 frames

[<ipython-input-9-30b5223765f1>](https://localhost:8080/#) in predict(input_img)
    301 def predict(input_img):
    302   # handle multi-stage outputs, obtain the last scale output of last stage
--> 303   return model.apply({'params': flax.core.freeze(params)}, input_img)
    304 
    305 

    [... skipping hidden 6 frame]

[/content/maxim/maxim/maxim/maxim/models/maxim.py](https://localhost:8080/#) in __call__(self, x, train)
    669   def __call__(self, x: jnp.ndarray, *, train: bool = False) -> Any:
    670 
--> 671     n, h, w, c = x.shape  # input image shape
    672     shortcuts = []
    673     shortcuts.append(x)

AttributeError: 'str' object has no attribute 'shape'
darkzbaron commented 10 months ago

That code works for me, hope this help!

`

images = ["../gdrive/My Drive/maxim/input/"+i for i in os.listdir("../gdrive/My Drive/maxim/input/") if i.endswith(('jpeg', 'png', 'jpg',"PNG","JPEG","JPG"))]

from io import BytesIO import time import cv2 input_dir = '/content/folder' model_action = '' images = [input_dir + "/" +i for i in os.listdir(inputdir) if i.endswith(('jpeg', 'png', 'jpg',"PNG","JPEG","JPG"))] print(images) = [resize(path, 1920, save_image=True) for path in images] # Resize Images to 1920 as the max dimension's size else it'll blow the GPU / CPU memory

for path in images: with open(path, 'rb') as f: image_bytes = BytesIO(f.read())

start = time.time() input_img, height, width, height_even, width_even = pre_process(image_bytes) preds = predict(input_img)

result = post_process(preds, height, width, height_even, width_even) img_ext = path.split('.')[-1] out_img = path.replace("." + img_ext,"_maxim." + img_ext)

result.save(input_dir+path.split('/')[-1])

cv2.imwrite(out_img, result) elapsed = time.time() - start minutes, seconds = divmod(elapsed, 60) # convert to minutes and remaining secondsv print("Time elapsed " + str(minutes) + "minutes and "+ str(seconds)) `