Views
No views yet
model folder, that you can dowload from this repository as well as on GitHub, present in your project folder.1import torch
2import torchvision.transforms as T
3from PIL import Image
4from huggingface_hub import hf_hub_download
5from model.MIRNet.model import MIRNet
6
7device = (
8 torch.device("cuda")
9 if torch.cuda.is_available()
10 else torch.device("mps")
11 if torch.backends.mps.is_available()
12 else torch.device("cpu")
13)
14
15# Download the model weights from the Hugging Face Hub
16model_path = hf_hub_download(
17 repo_id="dblasko/mirnet-low-light-img-enhancement", filename="mirnet_finetuned.pth"
18)
19
20# Load the model
21model = MIRNet().to(device)
22model.load_state_dict(torch.load(model_path, map_location=device)["model_state_dict"])
23
24# Use the model, for example for inference on an image
25model.eval()
26with torch.no_grad():
27 img = Image.open("image_path.png").convert("RGB")
28 img_tensor = T.Compose(
29 [
30 T.Resize(400), # Adjust image resizing depending on hardware
31 T.ToTensor(),
32 T.Normalize([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]),
33 ]
34 )(img).unsqueeze(0)
35 img_tensor = img_tensor.to(device)
36
37 if img_tensor.shape[2] % 8 != 0:
38 img_tensor = img_tensor[:, :, : -(img_tensor.shape[2] % 8), :]
39 if img_tensor.shape[3] % 8 != 0:
40 img_tensor = img_tensor[:, :, :, : -(img_tensor.shape[3] % 8)]
41
42 output = model(img_tensor)
43